diff --git a/.bazelrc b/.bazelrc index f9a0772778fd3..86c92328f4c68 100644 --- a/.bazelrc +++ b/.bazelrc @@ -17,9 +17,11 @@ build:no-tty --show_progress_rate_limit 10 build --define=cuda=true # rules_cuda configuration build --@rules_cuda//cuda:enable_cuda -build --@rules_cuda//cuda:cuda_targets=sm_52 +build --@rules_cuda//cuda:cuda_targets=sm_75 build --@rules_cuda//cuda:compiler=nvcc build --repo_env=CUDA_PATH=/usr/local/cuda +common --repo_env=PIP_INDEX_URL=https://pypi.org/simple/ +common --repo_env=PIP_TRUSTED_HOST=pypi.org # Configuration to build without GPU support build:cpu-only --define=cuda=false diff --git a/.ci/docker/almalinux/Dockerfile b/.ci/docker/almalinux/Dockerfile index b2bfb3e1212a7..d061353187c55 100644 --- a/.ci/docker/almalinux/Dockerfile +++ b/.ci/docker/almalinux/Dockerfile @@ -19,15 +19,12 @@ RUN git config --global --add safe.directory '*' ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH # cmake-3.18.4 from pip +# NS: Apr 1 2026 3.18.4 is gone, reported here https://github.com/scikit-build/cmake-python-distributions/issues/693 RUN yum install -y python3-pip && \ - python3 -mpip install cmake==3.18.4 && \ + python3 -mpip install cmake==3.18.4.post1 && \ ln -s /usr/local/bin/cmake /usr/bin/cmake3 RUN rm -rf /usr/local/cuda-* -FROM base as openssl -ADD ./common/install_openssl.sh install_openssl.sh -RUN bash ./install_openssl.sh && rm install_openssl.sh - FROM base as patchelf # Install patchelf ADD ./common/install_patchelf.sh install_patchelf.sh @@ -84,6 +81,8 @@ RUN yum -y update && \ yum -y install glibc-langpack-en && \ yum install -y sudo wget curl perl util-linux xz bzip2 git patch which perl zlib-devel openssl-devel yum-utils autoconf automake make gcc-toolset-${DEVTOOLSET_VERSION}-gcc gcc-toolset-${DEVTOOLSET_VERSION}-gcc-c++ gcc-toolset-${DEVTOOLSET_VERSION}-gcc-gfortran gcc-toolset-${DEVTOOLSET_VERSION}-gdb RUN git config --global --add safe.directory '*' +# All rocm clang cfg files load the same rocm.cfg, make sure it points to the right toolchain. +RUN echo "--gcc-toolchain=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr" >> /opt/rocm/llvm/bin/rocm.cfg ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH FROM rocm_base as rocm @@ -109,7 +108,6 @@ COPY --from=cuda13.2 /usr/local/cuda-13.2 /usr/local/cuda-13.2 # Final step FROM ${BASE_TARGET} as final ARG DEVTOOLSET_VERSION=13 -COPY --from=openssl /opt/openssl /opt/openssl COPY --from=patchelf /patchelf /usr/local/bin/patchelf COPY --from=conda /opt/conda /opt/conda diff --git a/.ci/docker/build.sh b/.ci/docker/build.sh index 23e9ef021c3ab..20c794a93dd4d 100755 --- a/.ci/docker/build.sh +++ b/.ci/docker/build.sh @@ -95,7 +95,6 @@ case "$tag" in CUDA_VERSION=12.4 ANACONDA_PYTHON_VERSION=3.10 GCC_VERSION=11 - VISION=yes KATEX=yes TRITON=yes ;; @@ -103,7 +102,6 @@ case "$tag" in CUDA_VERSION=12.8.1 ANACONDA_PYTHON_VERSION=3.10 GCC_VERSION=11 - VISION=yes KATEX=yes TRITON=yes INSTALL_MINGW=yes @@ -112,7 +110,6 @@ case "$tag" in CUDA_VERSION=13.0.2 ANACONDA_PYTHON_VERSION=3.10 GCC_VERSION=11 - VISION=yes KATEX=yes TRITON=yes ;; @@ -120,7 +117,6 @@ case "$tag" in CUDA_VERSION=13.0.2 ANACONDA_PYTHON_VERSION=3.10 GCC_VERSION=11 - VISION=yes KATEX=yes TRITON=yes INDUCTOR_BENCHMARKS=yes @@ -129,35 +125,32 @@ case "$tag" in CUDA_VERSION=13.0.2 ANACONDA_PYTHON_VERSION=3.12 GCC_VERSION=11 - VISION=yes KATEX=yes TRITON=yes ;; - pytorch-linux-jammy-py3-clang15-onnx) + pytorch-linux-jammy-py3.10-clang18) ANACONDA_PYTHON_VERSION=3.10 - CLANG_VERSION=15 - VISION=yes + CLANG_VERSION=18 + GCC_VERSION=11 + KATEX=yes + DOCS=yes ONNX=yes ;; - pytorch-linux-jammy-py3.10-clang15) - ANACONDA_PYTHON_VERSION=3.10 - CLANG_VERSION=15 - ;; - pytorch-linux-jammy-py3.11-clang15) + pytorch-linux-jammy-py3.11-clang18) ANACONDA_PYTHON_VERSION=3.11 - CLANG_VERSION=15 + CLANG_VERSION=18 ;; - pytorch-linux-jammy-py3.12-clang15) + pytorch-linux-jammy-py3.12-clang18) ANACONDA_PYTHON_VERSION=3.12 - CLANG_VERSION=15 + CLANG_VERSION=18 ;; - pytorch-linux-jammy-py3.13-clang15) + pytorch-linux-jammy-py3.13-clang18) ANACONDA_PYTHON_VERSION=3.13 - CLANG_VERSION=15 + CLANG_VERSION=18 ;; - pytorch-linux-jammy-py3.14-clang15) + pytorch-linux-jammy-py3.14-clang18) ANACONDA_PYTHON_VERSION=3.14 - CLANG_VERSION=15 + CLANG_VERSION=18 ;; pytorch-linux-jammy-rocm-n-py3 | pytorch-linux-jammy-rocm-n-py3-benchmarks | pytorch-linux-noble-rocm-n-py3) if [[ $tag =~ "jammy" ]]; then @@ -166,9 +159,7 @@ case "$tag" in ANACONDA_PYTHON_VERSION=3.12 fi GCC_VERSION=13 - VISION=yes ROCM_VERSION=7.2 - NINJA_VERSION=1.9.0 TRITON=yes KATEX=yes PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950;gfx1100" @@ -179,9 +170,7 @@ case "$tag" in pytorch-linux-noble-rocm-nightly-py3) ANACONDA_PYTHON_VERSION=3.12 GCC_VERSION=13 - VISION=yes ROCM_VERSION=nightly - NINJA_VERSION=1.9.0 TRITON=yes KATEX=yes PYTORCH_ROCM_ARCH="gfx942" @@ -189,23 +178,19 @@ case "$tag" in pytorch-linux-jammy-xpu-n-1-py3) ANACONDA_PYTHON_VERSION=3.10 GCC_VERSION=11 - VISION=yes XPU_VERSION=2025.2 XPU_DRIVER_TYPE=LTS - NINJA_VERSION=1.9.0 TRITON=yes ;; pytorch-linux-noble-xpu-n-py3 | pytorch-linux-noble-xpu-n-py3-client | pytorch-linux-noble-xpu-n-py3-inductor-benchmarks) ANACONDA_PYTHON_VERSION=3.10 GCC_VERSION=13 - VISION=yes XPU_VERSION=2025.3 if [[ $tag =~ "client" ]]; then XPU_DRIVER_TYPE=CLIENT else XPU_DRIVER_TYPE=LTS fi - NINJA_VERSION=1.9.0 TRITON=yes if [[ $tag =~ "benchmarks" ]]; then INDUCTOR_BENCHMARKS=yes @@ -214,36 +199,19 @@ case "$tag" in pytorch-linux-jammy-py3-gcc11-inductor-benchmarks) ANACONDA_PYTHON_VERSION=3.10 GCC_VERSION=11 - VISION=yes KATEX=yes - TRITON=yes DOCS=yes INDUCTOR_BENCHMARKS=yes ;; - pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang15) + pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang18) ANACONDA_PYTHON_VERSION=3.10 CUDA_VERSION=12.8.1 - CLANG_VERSION=15 - VISION=yes - TRITON=yes - ;; - pytorch-linux-jammy-py3-clang18-asan) - ANACONDA_PYTHON_VERSION=3.10 CLANG_VERSION=18 - VISION=yes - ;; - pytorch-linux-jammy-py3.10-gcc11) - ANACONDA_PYTHON_VERSION=3.10 - GCC_VERSION=11 - VISION=yes - KATEX=yes TRITON=yes - DOCS=yes - UNINSTALL_DILL=yes ;; - pytorch-linux-jammy-py3-clang15-executorch) + pytorch-linux-jammy-py3-clang18-executorch) ANACONDA_PYTHON_VERSION=3.10 - CLANG_VERSION=15 + CLANG_VERSION=18 EXECUTORCH=yes ;; pytorch-linux-jammy-py3.12-halide) @@ -288,21 +256,13 @@ case "$tag" in ANACONDA_PYTHON_VERSION=3.10 GCC_VERSION=13 ACL=yes - VISION=yes OPENBLAS=yes - # snadampal: skipping llvm src build install because the current version - # from pytorch/llvm:9.0.1 is x86 specific - SKIP_LLVM_SRC_BUILD_INSTALL=yes ;; pytorch-linux-jammy-aarch64-py3.10-gcc13-inductor-benchmarks) ANACONDA_PYTHON_VERSION=3.10 GCC_VERSION=13 ACL=yes - VISION=yes OPENBLAS=yes - # snadampal: skipping llvm src build install because the current version - # from pytorch/llvm:9.0.1 is x86 specific - SKIP_LLVM_SRC_BUILD_INSTALL=yes INDUCTOR_BENCHMARKS=yes ;; pytorch-linux-noble-riscv64-py3.12-gcc14) @@ -310,7 +270,6 @@ case "$tag" in ;; *) # Catch-all for builds that are not hardcoded. - VISION=yes echo "image '$image' did not match an existing build configuration" if [[ "$image" == *py* ]]; then extract_version_from_image_name py ANACONDA_PYTHON_VERSION @@ -327,14 +286,10 @@ case "$tag" in if [[ -z "$ROCM_VERSION" ]]; then extract_version_from_image_name rocm ROCM_VERSION fi - NINJA_VERSION=1.9.0 TRITON=yes # To ensure that any ROCm config will build using conda cmake # and thus have LAPACK/MKL enabled fi - if [[ "$image" == *centos7* ]]; then - NINJA_VERSION=1.10.2 - fi if [[ "$image" == *gcc* ]]; then extract_version_from_image_name gcc GCC_VERSION fi @@ -366,7 +321,6 @@ docker buildx build \ ${progress_flag} \ --build-arg "BUILD_ENVIRONMENT=${image}" \ --build-arg "LLVMDEV=${LLVMDEV:-}" \ - --build-arg "VISION=${VISION:-}" \ --build-arg "UBUNTU_VERSION=${UBUNTU_VERSION}" \ --build-arg "DEVTOOLSET_VERSION=${DEVTOOLSET_VERSION}" \ --build-arg "GLIBC_VERSION=${GLIBC_VERSION}" \ @@ -376,7 +330,6 @@ docker buildx build \ --build-arg "PYTHON_VERSION=${PYTHON_VERSION}" \ --build-arg "GCC_VERSION=${GCC_VERSION}" \ --build-arg "CUDA_VERSION=${CUDA_VERSION}" \ - --build-arg "NINJA_VERSION=${NINJA_VERSION:-}" \ --build-arg "KATEX=${KATEX:-}" \ --build-arg "ROCM_VERSION=${ROCM_VERSION:-}" \ --build-arg "PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}" \ @@ -392,11 +345,9 @@ docker buildx build \ --build-arg "TPU=${TPU}" \ --build-arg "XPU_VERSION=${XPU_VERSION}" \ --build-arg "XPU_DRIVER_TYPE=${XPU_DRIVER_TYPE}" \ - --build-arg "UNINSTALL_DILL=${UNINSTALL_DILL}" \ --build-arg "ACL=${ACL:-}" \ --build-arg "OPENBLAS=${OPENBLAS:-}" \ --build-arg "SKIP_SCCACHE_INSTALL=${SKIP_SCCACHE_INSTALL:-}" \ - --build-arg "SKIP_LLVM_SRC_BUILD_INSTALL=${SKIP_LLVM_SRC_BUILD_INSTALL:-}" \ --build-arg "INSTALL_MINGW=${INSTALL_MINGW:-}" \ -f $(dirname ${DOCKERFILE})/Dockerfile \ --load \ diff --git a/.ci/docker/ci_commit_pins/huggingface-requirements.txt b/.ci/docker/ci_commit_pins/huggingface-requirements.txt index 08538ff511057..51a16f10e0632 100644 --- a/.ci/docker/ci_commit_pins/huggingface-requirements.txt +++ b/.ci/docker/ci_commit_pins/huggingface-requirements.txt @@ -1,2 +1,2 @@ -transformers==5.2.0 +transformers==5.5.3 soxr==0.5.0 diff --git a/.ci/docker/ci_commit_pins/triton-xpu.txt b/.ci/docker/ci_commit_pins/triton-xpu.txt index e6c93e1b5432c..3bec45a60a706 100644 --- a/.ci/docker/ci_commit_pins/triton-xpu.txt +++ b/.ci/docker/ci_commit_pins/triton-xpu.txt @@ -1 +1 @@ -33f782efa9464adebb448ea1f1df1a64ec37ceb0 +21033c4e2be9b42c9e6ce7a39a70ead2aba279b4 diff --git a/.ci/docker/ci_commit_pins/triton.txt b/.ci/docker/ci_commit_pins/triton.txt index 23407b4d540c4..ab27e757c6e57 100644 --- a/.ci/docker/ci_commit_pins/triton.txt +++ b/.ci/docker/ci_commit_pins/triton.txt @@ -1 +1 @@ -9844da955a9db14ec69c9aac828ee9803085e288 +b4e20bbe55617cc798b986c2555a2bc7b303c737 diff --git a/.ci/docker/common/install_amdsmi.sh b/.ci/docker/common/install_amdsmi.sh index 8e0ee620da679..759e2bababe25 100644 --- a/.ci/docker/common/install_amdsmi.sh +++ b/.ci/docker/common/install_amdsmi.sh @@ -7,7 +7,7 @@ source /etc/rocm_env.sh # For theRock nightly, amd_smi may already be installed or in a different location if [ -d "${ROCM_PATH}/share/amd_smi" ]; then echo "Installing amdsmi from: ${ROCM_PATH}/share/amd_smi" - cd ${ROCM_PATH}/share/amd_smi && pip install . + cd ${ROCM_PATH}/share/amd_smi && python3 -m pip install . else echo "AMD SMI not found at ${ROCM_PATH}/share/amd_smi - skipping (may already be installed via pip)" fi diff --git a/.ci/docker/common/install_base.sh b/.ci/docker/common/install_base.sh index 7d8ae247d7a0b..27d9ca0601779 100755 --- a/.ci/docker/common/install_base.sh +++ b/.ci/docker/common/install_base.sh @@ -11,36 +11,26 @@ install_ubuntu() { # "$UBUNTU_VERSION" == "18.04" if [[ "$UBUNTU_VERSION" == "20.04"* ]]; then cmake3="cmake=3.16*" - maybe_libiomp_dev="" elif [[ "$UBUNTU_VERSION" == "22.04"* ]]; then cmake3="cmake=3.22*" - maybe_libiomp_dev="" elif [[ "$UBUNTU_VERSION" == "24.04"* ]]; then cmake3="cmake=3.28*" - maybe_libiomp_dev="" else - cmake3="cmake=3.5*" - maybe_libiomp_dev="libiomp-dev" - fi - - if [[ "$CLANG_VERSION" == 15 ]]; then - maybe_libomp_dev="libomp-15-dev" - elif [[ "$CLANG_VERSION" == 12 ]]; then - maybe_libomp_dev="libomp-12-dev" - elif [[ "$CLANG_VERSION" == 10 ]]; then - maybe_libomp_dev="libomp-10-dev" - else - maybe_libomp_dev="" + echo "Unknown Ubuntu version $UBUNTU_VERSION" + exit 1 fi # Install common dependencies apt-get update + # Install prerequisites for add-apt-repository (needs gpg-agent for PPA key import) + apt-get install -y --no-install-recommends software-properties-common gpg-agent + # Add git-core PPA for a newer version of git + add-apt-repository ppa:git-core/ppa -y + apt-get update # TODO: Some of these may not be necessary - ccache_deps="asciidoc docbook-xml docbook-xsl xsltproc" deploy_deps="libffi-dev libbz2-dev libreadline-dev libncurses5-dev libncursesw5-dev libgdbm-dev libsqlite3-dev uuid-dev tk-dev" numpy_deps="gfortran" apt-get install -y --no-install-recommends \ - $ccache_deps \ $numpy_deps \ ${deploy_deps} \ ${cmake3} \ @@ -53,14 +43,13 @@ install_ubuntu() { git \ libatlas-base-dev \ libc6-dbg \ - ${maybe_libiomp_dev} \ libyaml-dev \ libz-dev \ libjemalloc2 \ + libgl1 \ libjpeg-dev \ libasound2-dev \ libsndfile-dev \ - ${maybe_libomp_dev} \ software-properties-common \ wget \ sudo \ @@ -71,7 +60,9 @@ install_ubuntu() { unzip \ gpg-agent \ gdb \ - bc + bc \ + zip \ + valgrind # Should resolve issues related to various apt package repository cert issues # see: https://github.com/pytorch/pytorch/issues/65931 @@ -82,70 +73,14 @@ install_ubuntu() { rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* } -install_centos() { - # Need EPEL for many packages we depend on. - # See http://fedoraproject.org/wiki/EPEL - yum --enablerepo=extras install -y epel-release - - ccache_deps="asciidoc docbook-dtds docbook-style-xsl libxslt" - numpy_deps="gcc-gfortran" - yum install -y \ - $ccache_deps \ - $numpy_deps \ - autoconf \ - automake \ - bzip2 \ - cmake \ - cmake3 \ - curl \ - gcc \ - gcc-c++ \ - gflags-devel \ - git \ - glibc-devel \ - glibc-headers \ - glog-devel \ - libstdc++-devel \ - libsndfile-devel \ - make \ - opencv-devel \ - sudo \ - wget \ - vim \ - unzip \ - gdb - - # Cleanup - yum clean all - rm -rf /var/cache/yum - rm -rf /var/lib/yum/yumdb - rm -rf /var/lib/yum/history -} - # Install base packages depending on the base OS ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"') case "$ID" in ubuntu) install_ubuntu ;; - centos) - install_centos - ;; *) echo "Unable to determine OS..." exit 1 ;; esac - -# Install Valgrind separately since the apt-get version is too old. -mkdir valgrind_build && cd valgrind_build -VALGRIND_VERSION=3.20.0 -wget https://ossci-linux.s3.amazonaws.com/valgrind-${VALGRIND_VERSION}.tar.bz2 -tar -xjf valgrind-${VALGRIND_VERSION}.tar.bz2 -cd valgrind-${VALGRIND_VERSION} -./configure --prefix=/usr/local -make -j$[$(nproc) - 2] -sudo make install -cd ../../ -rm -rf valgrind_build -alias valgrind="/usr/local/bin/valgrind" diff --git a/.ci/docker/common/install_cache.sh b/.ci/docker/common/install_cache.sh index 040a31fc379d0..837ba50098029 100644 --- a/.ci/docker/common/install_cache.sh +++ b/.ci/docker/common/install_cache.sh @@ -18,7 +18,8 @@ install_ubuntu() { cp target/release/sccache-dist /opt/cache/bin echo "Cleaning up" cd .. - rm -rf sccache .cargo + rm -rf sccache + rustup self uninstall -y apt-get remove -y pkg-config libssl-dev apt-get autoclean && apt-get clean diff --git a/.ci/docker/common/install_clang.sh b/.ci/docker/common/install_clang.sh index 93daeee919b3d..20eea1ee157fe 100755 --- a/.ci/docker/common/install_clang.sh +++ b/.ci/docker/common/install_clang.sh @@ -15,7 +15,7 @@ if [ -n "$CLANG_VERSION" ]; then sudo apt-get update if [[ $CLANG_VERSION -ge 18 ]]; then - apt-get install -y libomp-${CLANG_VERSION}-dev libclang-rt-${CLANG_VERSION}-dev clang-"$CLANG_VERSION" llvm-"$CLANG_VERSION" + apt-get install -y --no-install-recommends libomp-${CLANG_VERSION}-dev libclang-rt-${CLANG_VERSION}-dev clang-"$CLANG_VERSION" llvm-"$CLANG_VERSION" else apt-get install -y --no-install-recommends clang-"$CLANG_VERSION" llvm-"$CLANG_VERSION" fi diff --git a/.ci/docker/common/install_conda.sh b/.ci/docker/common/install_conda.sh index a7e2536cf828d..547eec0d401f5 100755 --- a/.ci/docker/common/install_conda.sh +++ b/.ci/docker/common/install_conda.sh @@ -96,5 +96,8 @@ if [ -n "$ANACONDA_PYTHON_VERSION" ]; then pip_install -r /opt/conda/requirements-docs.txt fi + # Clean conda package cache + as_jenkins conda clean -ya + popd fi diff --git a/.ci/docker/common/install_conda_docker.sh b/.ci/docker/common/install_conda_docker.sh index dc377075750ac..9665a799b6ff5 100755 --- a/.ci/docker/common/install_conda_docker.sh +++ b/.ci/docker/common/install_conda_docker.sh @@ -13,8 +13,8 @@ rm $(basename "$MINICONDA_URL") export PATH=/opt/conda/bin:$PATH # See https://github.com/pytorch/builder/issues/1473 # Pin conda to 23.5.2 as it's the last one compatible with openssl-1.1.1 -conda install -y conda=23.5.2 conda-build anaconda-client git ninja +conda install -y conda=23.5.2 conda-build anaconda-client git # The cmake version here needs to match with the minimum version of cmake -# supported by PyTorch (3.18). There is only 3.18.2 on anaconda -/opt/conda/bin/pip3 install cmake==3.18.2 +# supported by PyTorch (3.18). +/opt/conda/bin/pip3 install cmake==3.18.4.post1 ninja conda remove -y --force patchelf diff --git a/.ci/docker/common/install_devtoolset.sh b/.ci/docker/common/install_devtoolset.sh deleted file mode 100755 index bdae637598138..0000000000000 --- a/.ci/docker/common/install_devtoolset.sh +++ /dev/null @@ -1,10 +0,0 @@ -#!/bin/bash - -set -ex - -[ -n "$DEVTOOLSET_VERSION" ] - -yum install -y centos-release-scl -yum install -y devtoolset-$DEVTOOLSET_VERSION - -echo "source scl_source enable devtoolset-$DEVTOOLSET_VERSION" > "/etc/profile.d/devtoolset-$DEVTOOLSET_VERSION.sh" diff --git a/.ci/docker/common/install_docs_reqs.sh b/.ci/docker/common/install_docs_reqs.sh index c907145f2ec62..c06160373a05e 100644 --- a/.ci/docker/common/install_docs_reqs.sh +++ b/.ci/docker/common/install_docs_reqs.sh @@ -17,7 +17,7 @@ if [ -n "$KATEX" ]; then apt-get install -y --no-install-recommends yarn yarn global add katex --prefix /usr/local - sudo apt-get -y install doxygen + sudo apt-get -y install doxygen lcov apt-get autoclean && apt-get clean rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* diff --git a/.ci/docker/common/install_glibc.sh b/.ci/docker/common/install_glibc.sh deleted file mode 100755 index c98791e2bf85b..0000000000000 --- a/.ci/docker/common/install_glibc.sh +++ /dev/null @@ -1,34 +0,0 @@ -#!/bin/bash - -set -ex - -[ -n "$GLIBC_VERSION" ] -if [[ -n "$CENTOS_VERSION" ]]; then - [ -n "$DEVTOOLSET_VERSION" ] -fi - -yum install -y wget sed - -mkdir -p /packages && cd /packages -wget -q http://ftp.gnu.org/gnu/glibc/glibc-$GLIBC_VERSION.tar.gz -tar xzf glibc-$GLIBC_VERSION.tar.gz -if [[ "$GLIBC_VERSION" == "2.26" ]]; then - cd glibc-$GLIBC_VERSION - sed -i 's/$name ne "nss_test1"/$name ne "nss_test1" \&\& $name ne "nss_test2"/' scripts/test-installation.pl - cd .. -fi -mkdir -p glibc-$GLIBC_VERSION-build && cd glibc-$GLIBC_VERSION-build - -if [[ -n "$CENTOS_VERSION" ]]; then - export PATH=/opt/rh/devtoolset-$DEVTOOLSET_VERSION/root/usr/bin:$PATH -fi - -../glibc-$GLIBC_VERSION/configure --prefix=/usr CFLAGS='-Wno-stringop-truncation -Wno-format-overflow -Wno-restrict -Wno-format-truncation -g -O2' -make -j$(nproc) -make install - -# Cleanup -rm -rf /packages -rm -rf /var/cache/yum/* -rm -rf /var/lib/rpm/__db.* -yum clean all diff --git a/.ci/docker/common/install_inductor_benchmark_deps.sh b/.ci/docker/common/install_inductor_benchmark_deps.sh index 674b141efcfb2..c54b8a44f0632 100644 --- a/.ci/docker/common/install_inductor_benchmark_deps.sh +++ b/.ci/docker/common/install_inductor_benchmark_deps.sh @@ -18,18 +18,16 @@ function install_timm() { function install_torchbench() { local commit commit=$(get_pinned_commit torchbench) - git clone https://github.com/pytorch/benchmark torchbench + mkdir torchbench && chown jenkins torchbench + as_jenkins git clone https://github.com/pytorch/benchmark torchbench pushd torchbench - git checkout "$commit" + as_jenkins git checkout "$commit" - python install.py --continue_on_fail + conda_run python install.py --continue_on_fail echo "Print all dependencies after TorchBench is installed" - python -mpip freeze + conda_run python -mpip freeze popd - - chown -R jenkins torchbench - chown -R jenkins /opt/conda } # Pango is needed for weasyprint which is needed for doctr diff --git a/.ci/docker/common/install_miopen.sh b/.ci/docker/common/install_miopen.sh index 3dbc67b90abaf..039458add8406 100644 --- a/.ci/docker/common/install_miopen.sh +++ b/.ci/docker/common/install_miopen.sh @@ -16,7 +16,7 @@ case "$ID" in ubuntu) IS_UBUNTU=1 ;; - centos|almalinux) + almalinux) IS_UBUNTU=0 ;; *) diff --git a/.ci/docker/common/install_ninja.sh b/.ci/docker/common/install_ninja.sh deleted file mode 100644 index fa380722bdc2f..0000000000000 --- a/.ci/docker/common/install_ninja.sh +++ /dev/null @@ -1,18 +0,0 @@ -#!/bin/bash - -set -ex - -[ -n "$NINJA_VERSION" ] - -arch=$(uname -m) -if [ "$arch" == "aarch64" ]; then - url="https://github.com/ninja-build/ninja/releases/download/v${NINJA_VERSION}/ninja-linux-aarch64.zip" -else - url="https://github.com/ninja-build/ninja/releases/download/v${NINJA_VERSION}/ninja-linux.zip" -fi - -pushd /tmp -wget --no-verbose --output-document=ninja-linux.zip "$url" -unzip ninja-linux.zip -d /usr/local/bin -rm -f ninja-linux.zip -popd \ No newline at end of file diff --git a/.ci/docker/common/install_onnx.sh b/.ci/docker/common/install_onnx.sh index 36ce5b11d9135..eacc7bbcad157 100755 --- a/.ci/docker/common/install_onnx.sh +++ b/.ci/docker/common/install_onnx.sh @@ -11,15 +11,11 @@ retry () { # ONNXRuntime should be installed before installing # onnx-weekly. Otherwise, onnx-weekly could be # overwritten by onnx. +# Note: parameterized, pytest-subtests, tabulate, packaging are already +# installed via requirements-ci.txt pip_install \ - parameterized==0.8.1 \ - pytest-cov==4.0.0 \ - pytest-subtests==0.10.0 \ - tabulate==0.9.0 \ - transformers==4.36.2 - -pip_install coloredlogs packaging -pip_install onnxruntime==1.23.1 + transformers==4.36.2 \ + onnxruntime==1.23.1 # Cache the transformers model to be used later by ONNX tests. We need to run the transformers # package to download the model. By default, the model is cached at ~/.cache/huggingface/hub/ @@ -34,4 +30,5 @@ conda_run python "${IMPORT_SCRIPT_FILENAME}" # Cleaning up conda_run pip uninstall -y torch +conda_run pip cache purge rm "${IMPORT_SCRIPT_FILENAME}" || true diff --git a/.ci/docker/common/install_openssl.sh b/.ci/docker/common/install_openssl.sh deleted file mode 100644 index c73c9c333c002..0000000000000 --- a/.ci/docker/common/install_openssl.sh +++ /dev/null @@ -1,17 +0,0 @@ -#!/bin/bash - -set -ex - -OPENSSL=openssl-1.1.1k - -wget -q -O "${OPENSSL}.tar.gz" "https://ossci-linux.s3.amazonaws.com/${OPENSSL}.tar.gz" -tar xf "${OPENSSL}.tar.gz" -cd "${OPENSSL}" -./config --prefix=/opt/openssl -d '-Wl,--enable-new-dtags,-rpath,$(LIBRPATH)' -# NOTE: openssl install errors out when built with the -j option -NPROC=$[$(nproc) - 2] -make -j${NPROC}; make install_sw -# Link the ssl libraries to the /usr/lib folder. -sudo ln -s /opt/openssl/lib/lib* /usr/lib -cd .. -rm -rf "${OPENSSL}" diff --git a/.ci/docker/common/install_rocSHMEM.sh b/.ci/docker/common/install_rocSHMEM.sh new file mode 100644 index 0000000000000..ff59951f94875 --- /dev/null +++ b/.ci/docker/common/install_rocSHMEM.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# Script used only in CD pipeline to build and install rocSHMEM + +set -eou pipefail + +function do_install() { + ROCSHMEM_VERSION=ea5c137103f18a9aadd570d09d72e78ec52f0a3a + rocm_dir="${ROCM_HOME:-}" + if [[ -z "${rocm_dir}" && -f /etc/rocm_env.sh ]]; then + source /etc/rocm_env.sh + rocm_dir="${ROCM_HOME:-}" + fi + rocm_dir="${rocm_dir:-/opt/rocm}" + echo "install_rocSHMEM.sh: using ROCM install prefix ${rocm_dir}" + if [[ -f "${rocm_dir}/lib/librocshmem.a" ]]; then + echo "install_rocSHMEM.sh: librocshmem.a already present in ${rocm_dir}/lib, skipping build" + return + fi + ( + set -x + curr_dir=$(pwd) + tmp_dir=$(mktemp -d) + + git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems.git ${tmp_dir}/rocm-systems + cd ${tmp_dir}/rocm-systems + git sparse-checkout set --cone projects/rocshmem + git checkout ${ROCSHMEM_VERSION} + + cd ${tmp_dir}/rocm-systems/projects/rocshmem + mkdir build + cd build + INSTALL_PREFIX="${rocm_dir}" ../scripts/build_configs/all_backends + cd ${curr_dir} + + ) +} + +do_install diff --git a/.ci/docker/common/install_rocm.sh b/.ci/docker/common/install_rocm.sh index 8b673a23f9de5..b4f6d70af3676 100644 --- a/.ci/docker/common/install_rocm.sh +++ b/.ci/docker/common/install_rocm.sh @@ -26,87 +26,131 @@ install_ubuntu() { apt-get install -y libc++1 apt-get install -y libc++abi1 - # When ROCM_VERSION=nightly, install ROCm from TheRock nightly wheels + # When ROCM_VERSION=nightly, install ROCm from TheRock nightly tarballs + # Mirrors: https://github.com/ROCm/TheRock/blob/main/dockerfiles/install_rocm_tarball.sh if [[ "${ROCM_VERSION}" == "nightly" ]]; then - echo "install_rocm.sh: installing ROCm from TheRock nightly wheels" + apt-get install -y --no-install-recommends pkg-config - # Clean any previous ROCm installation in the base CI image. if [[ -d /opt/rocm ]]; then - echo "Removing existing /opt/rocm from base image" rm -rf /opt/rocm fi - # Determine theRock nightly URL based on GPU architecture - # Check BUILD_ENVIRONMENT or PYTORCH_ROCM_ARCH for the target GPU - if [[ -z "${THEROCK_NIGHTLY_INDEX_URL:-}" ]]; then + # Determine GPU family based on target architecture + AMDGPU_FAMILY="${THEROCK_AMDGPU_FAMILY:-}" + if [[ -z "${AMDGPU_FAMILY}" ]]; then if [[ "${BUILD_ENVIRONMENT}" == *"gfx950"* ]] || [[ "${PYTORCH_ROCM_ARCH}" == *"gfx950"* ]]; then - # MI350 (gfx950) - THEROCK_NIGHTLY_INDEX_URL="https://rocm.nightlies.amd.com/v2/gfx950-dcgpu/" - echo "Detected gfx950 architecture - using MI350 theRock nightly repository" + AMDGPU_FAMILY="gfx950-dcgpu" else - # Default to MI300 (gfx942/gfx94X) - THEROCK_NIGHTLY_INDEX_URL="https://rocm.nightlies.amd.com/v2/gfx94X-dcgpu/" - echo "Using gfx94X (MI300) theRock nightly repository" + AMDGPU_FAMILY="gfx94X-dcgpu" fi fi - export THEROCK_NIGHTLY_INDEX_URL - echo "TheRock Index URL: ${THEROCK_NIGHTLY_INDEX_URL}" + # Auto-detect latest nightly version if not pinned + VERSION="${THEROCK_VERSION:-}" + if [[ -z "${VERSION}" ]]; then + VERSION=$(curl -fsSL "https://rocm.nightlies.amd.com/tarball/" \ + | grep -oP "therock-dist-linux-${AMDGPU_FAMILY}-\K[^\"]+(?=\.tar\.gz)" \ + | grep -v ADHOCBUILD \ + | sort -V \ + | tail -1) + if [[ -z "${VERSION}" ]]; then + echo "Error: Could not find a nightly tarball for ${AMDGPU_FAMILY}" + exit 1 + fi + fi + + # URL-encode '+' as '%2B' in VERSION (required for devreleases) + VERSION_ENCODED="${VERSION//+/%2B}" + + TARBALL_URL="https://rocm.nightlies.amd.com/tarball/therock-dist-linux-${AMDGPU_FAMILY}-${VERSION_ENCODED}.tar.gz" + + echo "==============================================" + echo "ROCm Tarball Installation" + echo "==============================================" + echo "Version: ${VERSION}" + echo "AMDGPU Family: ${AMDGPU_FAMILY}" + echo "Tarball URL: ${TARBALL_URL}" + echo "==============================================" + + # Download tarball + TARBALL_FILE="/tmp/rocm-tarball.tar.gz" + + echo "Downloading tarball..." + curl -fsSL -o "$TARBALL_FILE" "$TARBALL_URL" || { + echo "Error: Failed to download tarball from $TARBALL_URL" + exit 1 + } + + # Verify download + if [ ! -f "$TARBALL_FILE" ] || [ ! -s "$TARBALL_FILE" ]; then + echo "Error: Downloaded file is empty or does not exist" + exit 1 + fi + + # Install directory is fixed to /opt/rocm-{VERSION} + ROCM_INSTALL_DIR="/opt/rocm-${VERSION}" - python3 -m pip install \ - --index-url "${THEROCK_NIGHTLY_INDEX_URL}" \ - "rocm[libraries,devel]" + # Extract tarball to versioned directory + echo "Extracting tarball to ${ROCM_INSTALL_DIR}..." + mkdir -p "$ROCM_INSTALL_DIR" + tar -xzf "$TARBALL_FILE" -C "$ROCM_INSTALL_DIR" - # Use the rocm-sdk CLI helper to populate environment defaults - ROCM_HOME="$(rocm-sdk path --root)" - ROCM_BIN="$(rocm-sdk path --bin)" - ROCM_CMAKE_PREFIX="$(rocm-sdk path --cmake)" + # Clean up downloaded file + rm -f "$TARBALL_FILE" + echo "Tarball extracted and cleaned up" - echo "ROCM_HOME=${ROCM_HOME}" - echo "ROCM_BIN=${ROCM_BIN}" - echo "ROCM_CMAKE_PREFIX=${ROCM_CMAKE_PREFIX}" + # Create symlink /opt/rocm -> /opt/rocm-{VERSION} for compatibility + ln -sfn "$ROCM_INSTALL_DIR" /opt/rocm + echo "Created symlink: /opt/rocm -> $ROCM_INSTALL_DIR" - export ROCM_HOME - export ROCM_PATH="${ROCM_HOME}" - export PATH="${ROCM_BIN}:${PATH}" - export CMAKE_PREFIX_PATH="${ROCM_CMAKE_PREFIX}:${CMAKE_PREFIX_PATH:-}" + # Verify bin and lib folder exists after extraction + echo "Verifying installation..." + for dir in bin clients include lib libexec share; do + if [ ! -d "$ROCM_INSTALL_DIR/$dir" ]; then + echo "Error: ROCm $dir directory not found" + exit 1 + fi + echo "ROCm $dir found in $ROCM_INSTALL_DIR/$dir" + done - # theRock bundles system dependencies like libdrm, liblzma in rocm_sysdeps - ROCM_SYSDEPS="${ROCM_HOME}/lib/rocm_sysdeps" - ROCM_SYSDEPS_INCLUDE="${ROCM_SYSDEPS}/include" - ROCM_SYSDEPS_PKGCONFIG="${ROCM_SYSDEPS}/lib/pkgconfig" + echo "==============================================" + echo "ROCm installed successfully to $ROCM_INSTALL_DIR" + echo "ROCM_PATH=$ROCM_INSTALL_DIR" + echo "PATH should include: $ROCM_INSTALL_DIR/bin" + echo "==============================================" - # Write environment to file that can be sourced by CI scripts and users + # Write environment file (sourced by CI scripts and interactive shells) cat > /etc/rocm_env.sh << ROCM_ENV # ROCm paths -export ROCM_PATH="${ROCM_HOME}" -export ROCM_HOME="${ROCM_HOME}" -export ROCM_SOURCE_DIR="${ROCM_HOME}" -export ROCM_BIN="${ROCM_BIN}" -export ROCM_CMAKE="${ROCM_CMAKE_PREFIX}" -export PATH="${ROCM_BIN}:\${PATH}" -export CMAKE_PREFIX_PATH="${ROCM_CMAKE_PREFIX}:\${CMAKE_PREFIX_PATH:-}" -# Device library paths -export HIP_DEVICE_LIB_PATH="${ROCM_HOME}/lib/llvm/amdgcn/bitcode" -export ROCM_DEVICE_LIB_PATH="${ROCM_HOME}/lib/llvm/amdgcn/bitcode" -# theRock system dependencies -export ROCM_SYSDEPS_INCLUDE="${ROCM_SYSDEPS_INCLUDE}" -export CPLUS_INCLUDE_PATH="${ROCM_SYSDEPS_INCLUDE}:\${CPLUS_INCLUDE_PATH:-}" -export C_INCLUDE_PATH="${ROCM_SYSDEPS_INCLUDE}:\${C_INCLUDE_PATH:-}" -export PKG_CONFIG_PATH="${ROCM_SYSDEPS_PKGCONFIG}:\${PKG_CONFIG_PATH:-}" -export LD_LIBRARY_PATH="${ROCM_SYSDEPS}/lib:\${LD_LIBRARY_PATH:-}" -export LIBRARY_PATH="${ROCM_SYSDEPS}/lib:\${LIBRARY_PATH:-}" -export MAGMA_HOME="${ROCM_HOME}/magma" +export ROCM_PATH=/opt/rocm +export ROCM_HOME=/opt/rocm +export ROCM_SOURCE_DIR=/opt/rocm +export ROCM_BIN=/opt/rocm/bin +export ROCM_CMAKE=/opt/rocm +export PATH=/opt/rocm/bin:/opt/rocm/llvm/bin:\${PATH} +export LD_LIBRARY_PATH=/opt/rocm/lib:\${LD_LIBRARY_PATH:-} +# Sysdeps include paths (libdrm headers, etc.) +export CPLUS_INCLUDE_PATH=/opt/rocm/lib/rocm_sysdeps/include:\${CPLUS_INCLUDE_PATH:-} +export C_INCLUDE_PATH=/opt/rocm/lib/rocm_sysdeps/include:\${C_INCLUDE_PATH:-} +# Device library path +export HIP_DEVICE_LIB_PATH=/opt/rocm/amdgcn/bitcode +export MAGMA_HOME=/opt/rocm/magma +# Tarball bundles sysdeps (libdrm, liblzma, etc.); expose their libs and .pc files +if [ -d /opt/rocm/lib/rocm_sysdeps/lib ]; then + export LD_LIBRARY_PATH=/opt/rocm/lib/rocm_sysdeps/lib:\${LD_LIBRARY_PATH} + export PKG_CONFIG_PATH=/opt/rocm/lib/rocm_sysdeps/lib/pkgconfig:\${PKG_CONFIG_PATH:-} +fi # Disable MSLK for theRock nightly (not yet supported) export USE_MSLK=0 ROCM_ENV - # Append to bash.bashrc so interactive shells get the env vars echo "source /etc/rocm_env.sh" >> /etc/bash.bashrc - echo "install_rocm.sh: TheRock nightly ROCm install complete" - exit 0 - fi + # --- End of theRock nightly tarball installation --- + else + # ========================================================================= + # Non-nightly: install ROCm from repo.radeon.com apt packages + # ========================================================================= # Make sure rocm packages from repo.radeon.com have highest priority cat << EOF > /etc/apt/preferences.d/rocm-pin-600 @@ -120,6 +164,11 @@ EOF ROCM_VERSION="${ROCM_VERSION}.2" fi + # we want the patch version of 7.2 instead + if [[ $(ver $ROCM_VERSION) -eq $(ver 7.2) ]]; then + ROCM_VERSION="${ROCM_VERSION}.1" + fi + # Default url values rocm_baseurl="http://repo.radeon.com/rocm/apt/${ROCM_VERSION}" UBUNTU_VERSION_NAME=`cat /etc/os-release | grep UBUNTU_CODENAME | awk -F= '{print $2}'` @@ -154,29 +203,6 @@ EOF fi fi - # ROCm 7.2 needs a fix from procprof sdk that isn't available until 7.2.1 - if [[ $(ver $ROCM_VERSION) -eq $(ver 7.2) ]]; then - git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems.git - pushd rocm-systems/ - git sparse-checkout init --cone - git sparse-checkout set projects/rocprofiler-sdk shared/rocprofiler-compute - git checkout develop - git checkout rocm-7.2.0 - git config --global user.email "you@example.com" - git config --global user.name "Your Name" - git cherry-pick a71cc3cc88ed68b24c40cefec77d764053044862 - sudo apt install -y cmake libdw-dev libsqlite3-dev - cmake \ - -B rocprofiler-sdk-build \ - -DCMAKE_INSTALL_PREFIX=/opt/rocm \ - -DCMAKE_PREFIX_PATH=/opt/rocm \ - -DGPU_TARGETS="${PYTORCH_ROCM_ARCH}" \ - projects/rocprofiler-sdk - cmake --build rocprofiler-sdk-build --target all --parallel $(nproc) - cmake --build rocprofiler-sdk-build --target install - popd - fi - # ROCm 6.0 had a regression where journal_mode was enabled on the kdb files resulting in permission errors at runtime for kdb in /opt/rocm/share/miopen/db/*.kdb do @@ -217,7 +243,7 @@ EOF pip_install "git+https://github.com/rocm/composable_kernel@$ROCM_COMPOSABLE_KERNEL_VERSION" - # Write environment to file that can be sourced by CI scripts and users + # Write environment file (sourced by CI scripts and interactive shells) cat > /etc/rocm_env.sh << ROCM_ENV # ROCm paths export ROCM_PATH=/opt/rocm @@ -226,85 +252,18 @@ export ROCM_SOURCE_DIR=/opt/rocm export ROCM_BIN=/opt/rocm/bin export ROCM_CMAKE=/opt/rocm export PATH=/opt/rocm/bin:/opt/rocm/llvm/bin:\${PATH} -# Device library paths -export ROCM_DEVICE_LIB_PATH=/opt/rocm/amdgcn/bitcode +export LD_LIBRARY_PATH=/opt/rocm/lib:\${LD_LIBRARY_PATH:-} +# Device library path export HIP_DEVICE_LIB_PATH=/opt/rocm/amdgcn/bitcode export MAGMA_HOME=/opt/rocm/magma ROCM_ENV - # Append to bash.bashrc so interactive shells get the env vars echo "source /etc/rocm_env.sh" >> /etc/bash.bashrc # Cleanup apt-get autoclean && apt-get clean rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* -} - -install_centos() { - - yum update -y - yum install -y kmod - yum install -y wget - yum install -y openblas-devel - - yum install -y epel-release - yum install -y dkms kernel-headers-`uname -r` kernel-devel-`uname -r` - - # Add amdgpu repository - local amdgpu_baseurl - if [[ $OS_VERSION == 9 ]]; then - amdgpu_baseurl="https://repo.radeon.com/amdgpu/${ROCM_VERSION}/rhel/9.0/main/x86_64" - else - amdgpu_baseurl="https://repo.radeon.com/amdgpu/${ROCM_VERSION}/rhel/7.9/main/x86_64" - fi - echo "[AMDGPU]" > /etc/yum.repos.d/amdgpu.repo - echo "name=AMDGPU" >> /etc/yum.repos.d/amdgpu.repo - echo "baseurl=${amdgpu_baseurl}" >> /etc/yum.repos.d/amdgpu.repo - echo "enabled=1" >> /etc/yum.repos.d/amdgpu.repo - echo "gpgcheck=1" >> /etc/yum.repos.d/amdgpu.repo - echo "gpgkey=http://repo.radeon.com/rocm/rocm.gpg.key" >> /etc/yum.repos.d/amdgpu.repo - - local rocm_baseurl="http://repo.radeon.com/rocm/yum/${ROCM_VERSION}" - echo "[ROCm]" > /etc/yum.repos.d/rocm.repo - echo "name=ROCm" >> /etc/yum.repos.d/rocm.repo - echo "baseurl=${rocm_baseurl}" >> /etc/yum.repos.d/rocm.repo - echo "enabled=1" >> /etc/yum.repos.d/rocm.repo - echo "gpgcheck=1" >> /etc/yum.repos.d/rocm.repo - echo "gpgkey=http://repo.radeon.com/rocm/rocm.gpg.key" >> /etc/yum.repos.d/rocm.repo - - yum update -y - - yum install -y \ - rocm-dev \ - rocm-utils \ - rocm-libs \ - rccl \ - rocprofiler-dev \ - roctracer-dev \ - amd-smi-lib - - # precompiled miopen kernels; search for all unversioned packages - # if search fails it will abort this script; use true to avoid case where search fails - MIOPENHIPGFX=$(yum -q search miopen-hip-gfx | grep miopen-hip-gfx | awk '{print $1}'| grep -F kdb. || true) - if [[ "x${MIOPENHIPGFX}" = x ]]; then - echo "miopen-hip-gfx package not available" && exit 1 - else - yum install -y ${MIOPENHIPGFX} - fi - - # ROCm 6.0 had a regression where journal_mode was enabled on the kdb files resulting in permission errors at runtime - for kdb in /opt/rocm/share/miopen/db/*.kdb - do - sqlite3 $kdb "PRAGMA journal_mode=off; PRAGMA VACUUM;" - done - - pip_install "git+https://github.com/rocm/composable_kernel@$ROCM_COMPOSABLE_KERNEL_VERSION" - - # Cleanup - yum clean all - rm -rf /var/cache/yum - rm -rf /var/lib/yum/yumdb - rm -rf /var/lib/yum/history + fi } # Install Python packages depending on the base OS @@ -313,9 +272,6 @@ case "$ID" in ubuntu) install_ubuntu ;; - centos) - install_centos - ;; *) echo "Unable to determine OS..." exit 1 diff --git a/.ci/docker/common/install_rocm_drm.sh b/.ci/docker/common/install_rocm_drm.sh index c70f5880f2c5c..a6b0fe2c03924 100644 --- a/.ci/docker/common/install_rocm_drm.sh +++ b/.ci/docker/common/install_rocm_drm.sh @@ -14,7 +14,7 @@ case "$ID" in apt-get install -y libpciaccess-dev pkg-config apt-get clean ;; - centos|almalinux) + almalinux) yum install -y libpciaccess-devel pkgconfig ;; *) diff --git a/.ci/docker/common/install_torch_tpu.sh b/.ci/docker/common/install_torch_tpu.sh index 93326e44a6161..c4e4104edfe42 100644 --- a/.ci/docker/common/install_torch_tpu.sh +++ b/.ci/docker/common/install_torch_tpu.sh @@ -60,7 +60,7 @@ fetch_secret() { set +x fi - if ! gcloud secrets versions access latest --secret="torchtpu-readonly-key" --project="ml-velocity-actions-testing" > "temp_ssh_key"; then + if ! gcloud secrets versions access latest --secret="torchtpu-read-key" --project="ml-velocity-actions-testing" > "temp_ssh_key"; then echo "Error: Failed to fetch secret. Ensure you are authenticated with gcloud." # Restore xtrace if it was enabled, before exiting @@ -82,7 +82,7 @@ clone_repo() { # Use GIT_SSH_COMMAND to specify the key and disable strict host key checking for automation export GIT_SSH_COMMAND="ssh -i temp_ssh_key -o IdentitiesOnly=yes -o StrictHostKeyChecking=no" - if git clone --recursive "git@github.com:google-ml-infra/torch_tpu.git"; then + if git clone --recursive "git@github.com:google-pytorch/torch_tpu.git"; then echo "Repository cloned successfully." else echo "Error: Failed to clone repository." @@ -110,7 +110,7 @@ pull_torch_tpu() { # sleep 28800 # Debug sleep to connect to runner to streamline debugging, do not submit # 3. Configuration -TORCH_TPU_REPO="${TORCH_TPU_REPO:-https://github.com/google-ml-infra/torch_tpu.git}" +TORCH_TPU_REPO="${TORCH_TPU_REPO:-https://github.com/google-pytorch/torch_tpu.git}" TORCH_TPU_BRANCH="${TORCH_TPU_BRANCH:-main}" # Pin File Configuration diff --git a/.ci/docker/common/install_vision.sh b/.ci/docker/common/install_vision.sh deleted file mode 100755 index 78c445568ddcd..0000000000000 --- a/.ci/docker/common/install_vision.sh +++ /dev/null @@ -1,46 +0,0 @@ -#!/bin/bash - -set -ex - -install_ubuntu() { - apt-get update - apt-get install -y --no-install-recommends \ - libopencv-dev - - # Cleanup - apt-get autoclean && apt-get clean - rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* -} - -install_centos() { - # Need EPEL for many packages we depend on. - # See http://fedoraproject.org/wiki/EPEL - yum --enablerepo=extras install -y epel-release - - yum install -y \ - opencv-devel - - # Cleanup - yum clean all - rm -rf /var/cache/yum - rm -rf /var/lib/yum/yumdb - rm -rf /var/lib/yum/history -} - -# Install base packages depending on the base OS -ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"') -case "$ID" in - ubuntu) - install_ubuntu - ;; - centos) - install_centos - ;; - *) - echo "Unable to determine OS..." - exit 1 - ;; -esac - -# Cache vision models used by the test -source "$(dirname "${BASH_SOURCE[0]}")/cache_vision_models.sh" diff --git a/.ci/docker/manywheel/Dockerfile_2_28 b/.ci/docker/manywheel/Dockerfile_2_28 index 4055e6b872539..2ad5cd6498249 100644 --- a/.ci/docker/manywheel/Dockerfile_2_28 +++ b/.ci/docker/manywheel/Dockerfile_2_28 @@ -12,20 +12,12 @@ RUN yum install -y sudo wget curl perl util-linux xz bzip2 git patch which perl ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH -# cmake-3.18.4 from pip +# cmake-3.18.4.post1 from pip +# NS: Apr 1 2026 3.18.4 is gone, reported here https://github.com/scikit-build/cmake-python-distributions/issues/693 RUN yum install -y python3-pip && \ - python3 -mpip install cmake==3.18.4 && \ + python3 -mpip install cmake==3.18.4.post1 && \ ln -s /usr/local/bin/cmake /usr/bin/cmake3 -FROM base as openssl -# Install openssl (this must precede `build python` step) -# (In order to have a proper SSL module, Python is compiled -# against a recent openssl [see env vars above], which is linked -# statically. We delete openssl afterwards.) -ADD ./common/install_openssl.sh install_openssl.sh -RUN bash ./install_openssl.sh && rm install_openssl.sh - - FROM base as cuda ARG BASE_CUDA_VERSION=12.6 # Install CUDA @@ -95,7 +87,6 @@ RUN git config --global --add safe.directory "*" ENV SSL_CERT_FILE=/opt/_internal/certs.pem # Install LLVM version -COPY --from=openssl /opt/openssl /opt/openssl COPY --from=base /opt/python /opt/python COPY --from=base /usr/local/lib/ /usr/local/lib/ COPY --from=base /opt/_internal /opt/_internal @@ -127,9 +118,9 @@ RUN for cpython_version in "cp312-cp312" "cp313-cp313" "cp313-cp313t"; do \ ADD ./common/patch_libstdc.sh patch_libstdc.sh RUN bash ./patch_libstdc.sh && rm patch_libstdc.sh -# cmake-3.18.4 from pip; force in case cmake3 already exists +# cmake-3.18.4.post1 from pip; force in case cmake3 already exists RUN yum install -y python3-pip && \ - python3 -mpip install cmake==3.18.4 && \ + python3 -mpip install cmake==3.18.4.post1 && \ ln -sf /usr/local/bin/cmake /usr/bin/cmake3 FROM cpu_final as cuda_final @@ -144,7 +135,8 @@ ARG ROCM_VERSION=6.0 ARG PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH} ARG DEVTOOLSET_VERSION=13 -ENV LDFLAGS="-Wl,-rpath=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64 -Wl,-rpath=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib" +# All rocm clang cfg files load the same rocm.cfg, make sure it points to the right toolchain. +RUN echo "--gcc-toolchain=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr" >> /opt/rocm/llvm/bin/rocm.cfg # Somewhere in ROCm stack, we still use non-existing /opt/rocm/hip path, # below workaround helps avoid error ENV ROCM_PATH /opt/rocm @@ -160,6 +152,10 @@ RUN yum install -y libdrm-devel ENV MKLROOT /opt/intel ADD ./common/install_rocm_magma.sh install_rocm_magma.sh RUN bash ./install_rocm_magma.sh ${ROCM_VERSION} && rm install_rocm_magma.sh + +ADD ./common/install_rocSHMEM.sh install_rocSHMEM.sh +RUN bash ./install_rocSHMEM.sh ${ROCM_VERSION} && rm install_rocSHMEM.sh + ADD ./common/install_miopen.sh install_miopen.sh RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh diff --git a/.ci/docker/manywheel/Dockerfile_2_28_aarch64 b/.ci/docker/manywheel/Dockerfile_2_28_aarch64 index b5bf2ffc1c081..477e4221cb49f 100644 --- a/.ci/docker/manywheel/Dockerfile_2_28_aarch64 +++ b/.ci/docker/manywheel/Dockerfile_2_28_aarch64 @@ -39,12 +39,7 @@ RUN yum install -y \ gcc-toolset-${GCCTOOLSET_VERSION}-gcc-c++ \ gcc-toolset-${GCCTOOLSET_VERSION}-gcc-gfortran \ gcc-toolset-${GCCTOOLSET_VERSION}-gdb - -# (optional) Install non-default Ninja version -ARG NINJA_VERSION -COPY ./common/install_ninja.sh install_ninja.sh -RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi -RUN rm install_ninja.sh +RUN yum install -y --enablerepo=powertools ninja-build # Ensure the expected devtoolset is used ENV PATH=/opt/rh/gcc-toolset-${GCCTOOLSET_VERSION}/root/usr/bin:$PATH diff --git a/.ci/docker/manywheel/Dockerfile_cuda_aarch64 b/.ci/docker/manywheel/Dockerfile_cuda_aarch64 index 794a791b2721a..2a3f266fc413f 100644 --- a/.ci/docker/manywheel/Dockerfile_cuda_aarch64 +++ b/.ci/docker/manywheel/Dockerfile_cuda_aarch64 @@ -50,16 +50,7 @@ ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/op RUN git config --global --add safe.directory "*" -FROM base as openssl -# Install openssl (this must precede `build python` step) -# (In order to have a proper SSL module, Python is compiled -# against a recent openssl [see env vars above], which is linked -# statically. We delete openssl afterwards.) -ADD ./common/install_openssl.sh install_openssl.sh -RUN bash ./install_openssl.sh && rm install_openssl.sh -ENV SSL_CERT_FILE=/opt/_internal/certs.pem - -FROM openssl as final +FROM base as final FROM base as cuda ARG BASE_CUDA_VERSION diff --git a/.ci/docker/manywheel/build.sh b/.ci/docker/manywheel/build.sh index 60a7216ce5f7b..05e468d72dda9 100755 --- a/.ci/docker/manywheel/build.sh +++ b/.ci/docker/manywheel/build.sh @@ -40,7 +40,7 @@ case ${image} in manylinux2_28_aarch64-builder:cpu-aarch64) TARGET=final GPU_IMAGE=arm64v8/almalinux:8 - DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=13 --build-arg NINJA_VERSION=1.12.1" + DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=13" MANY_LINUX_VERSION="2_28_aarch64" ;; manylinuxs390x-builder:cpu-s390x) @@ -75,6 +75,10 @@ case ${image} in DOCKERFILE_SUFFIX="_cuda_aarch64" ;; manylinux2_28-builder:rocm*) + # we want the patch version of 7.2 instead + if [[ "$GPU_ARCH_VERSION" == *"7.2"* ]]; then + GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.1" + fi # we want the patch version of 7.1 instead if [[ "$GPU_ARCH_VERSION" == *"7.1"* ]]; then GPU_ARCH_VERSION="${GPU_ARCH_VERSION}.1" diff --git a/.ci/docker/requirements-ci.txt b/.ci/docker/requirements-ci.txt index 9eac1292c7d84..a8096959ee619 100644 --- a/.ci/docker/requirements-ci.txt +++ b/.ci/docker/requirements-ci.txt @@ -75,7 +75,8 @@ librosa==0.10.2 ; python_version == "3.12" and platform_machine != "s390x" #test that import: test_spectral_ops.py #librosa depends on numba; disable it for s390x while numba is disabled too -mkl==2024.2.0 ; platform_machine != "aarch64" and sys_platform != "darwin" +# Only mkl-static and mkl-include are needed; the mkl package contains +# dynamic libraries that are not discoverable by our build scripts. mkl-static==2024.2.0 ; platform_machine != "aarch64" and sys_platform != "darwin" mkl-include==2024.2.0 ; platform_machine != "aarch64" and sys_platform != "darwin" #Description: Intel oneAPI Math Kernel Library @@ -168,7 +169,7 @@ optree==0.17.0 ; python_version >= "3.14" #test_pointwise_ops.py, test_dtensor_ops.py, test_torchinductor.py, test_fx.py, #test_fake_tensor.py, test_mps.py -pillow==12.1.1 +pillow==12.2.0 #Description: Python Imaging Library fork #Pinned versions: 11.0.0 #test that import: @@ -223,7 +224,7 @@ xdoctest==1.3.0 #Pinned versions: 1.1.0 #test that import: -pygments==2.15.0 +pygments==2.20.0 #Description: support doctest highlighting #Pinned versions: 2.12.0 #test that import: the doctests @@ -344,7 +345,7 @@ sympy==1.13.3 #Pinned versions: #test that import: -onnx==1.20.0 +onnx==1.21.0 #Description: Required by the torch.onnx exporter #Pinned versions: #test that import: @@ -416,5 +417,5 @@ tqdm>=4.66.0 #Description: progress bar library required for dynamo benchmarks #test that import: benchmarks/dynamo/* -aiohttp==3.13.3 +aiohttp==3.13.4 #Description: required for torch.distributed.debug diff --git a/.ci/docker/requirements-docs.txt b/.ci/docker/requirements-docs.txt index 7e556a80e0025..38258e0189b72 100644 --- a/.ci/docker/requirements-docs.txt +++ b/.ci/docker/requirements-docs.txt @@ -6,13 +6,10 @@ pytorch_sphinx_theme2==0.4.6 #Description: This is needed to generate PyTorch docs #Pinned versions: 0.4.6 -# TODO: sphinxcontrib.katex 0.9.0 adds a local KaTeX server to speed up pre-rendering -# but it doesn't seem to work and hangs around idly. The initial thought that it is probably -# something related to Docker setup. We can investigate this later. - -sphinxcontrib.katex==0.8.6 +sphinxcontrib.katex==0.9.11 #Description: This is used to generate PyTorch docs -#Pinned versions: 0.8.6 +#Pinned versions: 0.9.11 (0.9.0+ uses a persistent KaTeX server instead of +# spawning a subprocess per math expression, ~20% faster writes) sphinxext-opengraph==0.9.1 #Description: This is used to generate PyTorch docs @@ -48,6 +45,10 @@ docutils==0.20 #Description: This is used to generate PyTorch C++ docs #Pinned versions: 0.20 +coverxygen==1.8.1 +#Description: This is used to measure C++ API doc coverage from Doxygen XML +#Pinned versions: 1.8.1 + bs4==0.0.1 #Description: This is used to generate PyTorch C++ docs #Pinned versions: 0.0.1 diff --git a/.ci/docker/triton_version.txt b/.ci/docker/triton_version.txt index 40c341bdcdbe8..7c69a55dbb185 100644 --- a/.ci/docker/triton_version.txt +++ b/.ci/docker/triton_version.txt @@ -1 +1 @@ -3.6.0 +3.7.0 diff --git a/.ci/docker/triton_xpu_version.txt b/.ci/docker/triton_xpu_version.txt index 7c69a55dbb185..a76ccff2a6e0d 100644 --- a/.ci/docker/triton_xpu_version.txt +++ b/.ci/docker/triton_xpu_version.txt @@ -1 +1 @@ -3.7.0 +3.7.1 diff --git a/.ci/docker/ubuntu-rocm/Dockerfile b/.ci/docker/ubuntu-rocm/Dockerfile index ff62e4a934c74..a3d697f1b27f2 100644 --- a/.ci/docker/ubuntu-rocm/Dockerfile +++ b/.ci/docker/ubuntu-rocm/Dockerfile @@ -13,6 +13,7 @@ ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH} # Install common dependencies (so that this step can be cached separately) COPY ./common/install_base.sh install_base.sh RUN bash ./install_base.sh && rm install_base.sh +RUN apt-get update && apt-get install -y --no-install-recommends libtbb-dev && rm -rf /var/lib/apt/lists/* # Install user COPY ./common/install_user.sh install_user.sh @@ -43,13 +44,6 @@ ARG CLANG_VERSION COPY ./common/install_clang.sh install_clang.sh RUN bash ./install_clang.sh && rm install_clang.sh -# (optional) Install vision packages like OpenCV -ARG VISION -COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./ -RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi -RUN rm install_vision.sh cache_vision_models.sh common_utils.sh -ENV INSTALLED_VISION ${VISION} - # Install rocm ARG ROCM_VERSION ENV ROCM_VERSION=${ROCM_VERSION} @@ -63,14 +57,27 @@ RUN rm -r ci_commit_pins COPY ./common/install_rocm_magma.sh install_rocm_magma.sh RUN if [ "${ROCM_VERSION}" != "nightly" ]; then bash ./install_rocm_magma.sh ${ROCM_VERSION}; fi RUN rm install_rocm_magma.sh +COPY ./common/install_rocSHMEM.sh install_rocSHMEM.sh +RUN bash ./install_rocSHMEM.sh ${ROCM_VERSION} +RUN rm install_rocSHMEM.sh ADD ./common/install_miopen.sh install_miopen.sh RUN if [ "${ROCM_VERSION}" != "nightly" ]; then bash ./install_miopen.sh ${ROCM_VERSION}; fi && rm install_miopen.sh ADD ./common/install_rocm_drm.sh install_rocm_drm.sh RUN if [ "${ROCM_VERSION}" != "nightly" ]; then bash ./install_rocm_drm.sh /usr ; fi && rm install_rocm_drm.sh -# ROCm environment variables are set in /etc/rocm_env.sh by install_rocm.sh -# and sourced via /etc/bash.bashrc for interactive shells. -# CI scripts should source /etc/rocm_env.sh directly. +# Default ROCm environment; /etc/rocm_env.sh (created by install_rocm.sh) may +# override these at runtime for different install methods (tarballs vs wheels). +ENV ROCM_PATH=/opt/rocm \ + ROCM_HOME=/opt/rocm \ + ROCM_SOURCE_DIR=/opt/rocm \ + ROCM_BIN=/opt/rocm/bin \ + ROCM_CMAKE=/opt/rocm \ + ROCM_DEVICE_LIB_PATH=/opt/rocm/amdgcn/bitcode \ + HIP_DEVICE_LIB_PATH=/opt/rocm/amdgcn/bitcode \ + MAGMA_HOME=/opt/rocm/magma +ENV PATH=/opt/rocm/bin:/opt/rocm/llvm/bin:$PATH +ENV LD_LIBRARY_PATH=/opt/rocm/lib:${LD_LIBRARY_PATH:-} + ENV LANG C.UTF-8 ENV LC_ALL C.UTF-8 @@ -79,11 +86,6 @@ COPY ./common/install_amdsmi.sh install_amdsmi.sh RUN bash ./install_amdsmi.sh RUN rm install_amdsmi.sh -COPY ./common/install_openssl.sh install_openssl.sh -ENV OPENSSL_ROOT_DIR /opt/openssl -RUN bash ./install_openssl.sh -ENV OPENSSL_DIR /opt/openssl - ARG INDUCTOR_BENCHMARKS ARG ANACONDA_PYTHON_VERSION ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION @@ -95,12 +97,6 @@ COPY ci_commit_pins/torchbench.txt torchbench.txt RUN if [ -n "${INDUCTOR_BENCHMARKS}" ]; then bash ./install_inductor_benchmark_deps.sh; fi RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface-requirements.txt torchbench.txt -# (optional) Install non-default Ninja version -ARG NINJA_VERSION -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 @@ -126,8 +122,5 @@ RUN rm install_openmpi.sh ARG BUILD_ENVIRONMENT ENV BUILD_ENVIRONMENT ${BUILD_ENVIRONMENT} -# Install LLVM dev version (Defined in the pytorch/builder github repository) -COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm - USER jenkins CMD ["bash"] diff --git a/.ci/docker/ubuntu-xpu/Dockerfile b/.ci/docker/ubuntu-xpu/Dockerfile index c61612882032d..3aba712da78d7 100644 --- a/.ci/docker/ubuntu-xpu/Dockerfile +++ b/.ci/docker/ubuntu-xpu/Dockerfile @@ -47,12 +47,6 @@ RUN bash ./install_gcc.sh && rm install_gcc.sh COPY ./common/install_lcov.sh install_lcov.sh RUN bash ./install_lcov.sh && rm install_lcov.sh -COPY ./common/install_openssl.sh install_openssl.sh -RUN bash ./install_openssl.sh -ENV OPENSSL_ROOT_DIR /opt/openssl -ENV OPENSSL_DIR /opt/openssl -RUN rm install_openssl.sh - ARG INDUCTOR_BENCHMARKS ARG ANACONDA_PYTHON_VERSION ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION @@ -80,19 +74,6 @@ COPY triton_xpu_version.txt triton_version.txt RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi RUN rm install_triton.sh common_utils.sh triton-xpu.txt triton_version.txt -# (optional) Install vision packages like OpenCV -ARG VISION -COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./ -RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi -RUN rm install_vision.sh cache_vision_models.sh common_utils.sh -ENV INSTALLED_VISION ${VISION} - -# (optional) Install non-default Ninja version -ARG NINJA_VERSION -COPY ./common/install_ninja.sh install_ninja.sh -RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi -RUN rm install_ninja.sh - # 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 @@ -102,8 +83,5 @@ RUN bash ./install_cache.sh && rm install_cache.sh ARG BUILD_ENVIRONMENT ENV BUILD_ENVIRONMENT ${BUILD_ENVIRONMENT} -# Install LLVM dev version (Defined in the pytorch/builder github repository) -COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm - USER jenkins CMD ["bash"] diff --git a/.ci/docker/ubuntu/Dockerfile b/.ci/docker/ubuntu/Dockerfile index 64ba3a86ddfdd..6c98d6639fb16 100644 --- a/.ci/docker/ubuntu/Dockerfile +++ b/.ci/docker/ubuntu/Dockerfile @@ -42,7 +42,6 @@ COPY ./common/install_conda.sh install_conda.sh COPY ./common/common_utils.sh common_utils.sh COPY ./common/install_magma_conda.sh install_magma_conda.sh RUN bash ./install_conda.sh && rm install_conda.sh install_magma_conda.sh common_utils.sh /opt/conda/requirements-ci.txt /opt/conda/requirements-docs.txt -RUN if [ -n "${UNINSTALL_DILL}" ]; then pip uninstall -y dill; fi # Install gcc ARG GCC_VERSION @@ -70,25 +69,6 @@ ENV NCCL_LIB_DIR="/usr/local/cuda/lib64/" ARG CUDA_VERSION -# (optional) Install vision packages like OpenCV -ARG VISION -COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./ -RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi -RUN rm install_vision.sh cache_vision_models.sh common_utils.sh -ENV INSTALLED_VISION ${VISION} - -# (optional) Install non-default Ninja version -ARG NINJA_VERSION -COPY ./common/install_ninja.sh install_ninja.sh -RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi -RUN rm install_ninja.sh - -COPY ./common/install_openssl.sh install_openssl.sh -RUN bash ./install_openssl.sh -ENV OPENSSL_ROOT_DIR /opt/openssl -ENV OPENSSL_DIR /opt/openssl -RUN rm install_openssl.sh - ARG INDUCTOR_BENCHMARKS COPY ./common/install_inductor_benchmark_deps.sh install_inductor_benchmark_deps.sh COPY ./common/common_utils.sh common_utils.sh @@ -197,11 +177,6 @@ RUN rm install_openmpi.sh ARG BUILD_ENVIRONMENT ENV BUILD_ENVIRONMENT ${BUILD_ENVIRONMENT} -# Install LLVM dev version (Defined in the pytorch/builder github repository) -ARG SKIP_LLVM_SRC_BUILD_INSTALL -COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm -RUN if [ -n "${SKIP_LLVM_SRC_BUILD_INSTALL}" ]; then set -eu; rm -rf /opt/llvm; fi - # AWS specific CUDA build guidance ENV TORCH_NVCC_FLAGS "-Xfatbin -compress-all" ENV CUDA_PATH /usr/local/cuda diff --git a/.ci/lumen_cli/pyproject.toml b/.ci/lumen_cli/pyproject.toml index b2ac379e34ab0..ce8cf59d99fda 100644 --- a/.ci/lumen_cli/pyproject.toml +++ b/.ci/lumen_cli/pyproject.toml @@ -6,7 +6,7 @@ dependencies = [ "GitPython==3.1.45", "docker==7.1.0", "pytest==7.3.2", - "uv==0.9.6" + "uv==0.11.6" ] [tool.setuptools] diff --git a/.ci/manywheel/build_common.sh b/.ci/manywheel/build_common.sh index d50bd623dace0..d76b75dea6850 100644 --- a/.ci/manywheel/build_common.sh +++ b/.ci/manywheel/build_common.sh @@ -91,11 +91,6 @@ export PYTORCH_BUILD_NUMBER=$build_number export CMAKE_LIBRARY_PATH="/opt/intel/lib:/lib:$CMAKE_LIBRARY_PATH" export CMAKE_INCLUDE_PATH="/opt/intel/include:$CMAKE_INCLUDE_PATH" -if [[ -e /opt/openssl ]]; then - export OPENSSL_ROOT_DIR=/opt/openssl - export CMAKE_INCLUDE_PATH="/opt/openssl/include":$CMAKE_INCLUDE_PATH -fi - mkdir -p /tmp/$WHEELHOUSE_DIR export PATCHELF_BIN=/usr/local/bin/patchelf @@ -118,6 +113,9 @@ retry pip install -qUr requirements-build.txt python setup.py clean retry pip install -qr requirements.txt case ${DESIRED_PYTHON} in + cp314*) + retry pip install -q --pre numpy==2.3.4 + ;; cp31*) retry pip install -q --pre numpy==2.1.0 ;; diff --git a/.ci/manywheel/build_libtorch.sh b/.ci/manywheel/build_libtorch.sh index d78fbd5c3ed36..852ecf7500604 100644 --- a/.ci/manywheel/build_libtorch.sh +++ b/.ci/manywheel/build_libtorch.sh @@ -59,12 +59,6 @@ export PYTORCH_BUILD_NUMBER=$build_number export CMAKE_LIBRARY_PATH="/opt/intel/lib:/lib:$CMAKE_LIBRARY_PATH" export CMAKE_INCLUDE_PATH="/opt/intel/include:$CMAKE_INCLUDE_PATH" -# set OPENSSL_ROOT_DIR=/opt/openssl if it exists -if [[ -e /opt/openssl ]]; then - export OPENSSL_ROOT_DIR=/opt/openssl - export CMAKE_INCLUDE_PATH="/opt/openssl/include":$CMAKE_INCLUDE_PATH -fi - # If given a python version like 3.6m or 2.7mu, convert this to the format we # expect. The binary CI jobs pass in python versions like this; they also only # ever pass one python version, so we assume that DESIRED_PYTHON is not a list diff --git a/.ci/manywheel/build_rocm.sh b/.ci/manywheel/build_rocm.sh index bac56746f4501..fa5724dca25a7 100755 --- a/.ci/manywheel/build_rocm.sh +++ b/.ci/manywheel/build_rocm.sh @@ -97,20 +97,13 @@ ROCM_SO_FILES=( "libhipblaslt.so" "libhipsparselt.so" "libhiprtc.so" + "librocprofiler-sdk.so" + "librocprofiler-register.so" + "libhsa-amd-aqlprofile64.so" + "librocm-core.so" + "librocroller.so" ) -if [[ $ROCM_INT -ge 60100 ]]; then - ROCM_SO_FILES+=("librocprofiler-register.so") -fi - -if [[ $ROCM_INT -ge 60200 ]]; then - ROCM_SO_FILES+=("librocm-core.so") -fi - -if [[ $ROCM_INT -ge 70000 ]]; then - ROCM_SO_FILES+=("librocroller.so") -fi - OS_NAME=`awk -F= '/^NAME/{print $2}' /etc/os-release` if [[ "$OS_NAME" == *"CentOS Linux"* || "$OS_NAME" == *"AlmaLinux"* ]]; then LIBGOMP_PATH="/usr/lib64/libgomp.so.1" @@ -121,61 +114,23 @@ if [[ "$OS_NAME" == *"CentOS Linux"* || "$OS_NAME" == *"AlmaLinux"* ]]; then else LIBTINFO_PATH="/usr/lib64/libtinfo.so.6" fi + LIBDW_PATH="/usr/lib64/libdw.so.1" LIBDRM_PATH="/opt/amdgpu/lib64/libdrm.so.2" LIBDRM_AMDGPU_PATH="/opt/amdgpu/lib64/libdrm_amdgpu.so.1" - if [[ $ROCM_INT -ge 60100 && $ROCM_INT -lt 60300 ]]; then - # Below libs are direct dependencies of libhipsolver - LIBSUITESPARSE_CONFIG_PATH="/lib64/libsuitesparseconfig.so.4" - if [[ "$OS_NAME" == *"CentOS Linux"* ]]; then - LIBCHOLMOD_PATH="/lib64/libcholmod.so.2" - # Below libs are direct dependencies of libsatlas - LIBGFORTRAN_PATH="/lib64/libgfortran.so.3" - else - LIBCHOLMOD_PATH="/lib64/libcholmod.so.3" - # Below libs are direct dependencies of libsatlas - LIBGFORTRAN_PATH="/lib64/libgfortran.so.5" - fi - # Below libs are direct dependencies of libcholmod - LIBAMD_PATH="/lib64/libamd.so.2" - LIBCAMD_PATH="/lib64/libcamd.so.2" - LIBCCOLAMD_PATH="/lib64/libccolamd.so.2" - LIBCOLAMD_PATH="/lib64/libcolamd.so.2" - LIBSATLAS_PATH="/lib64/atlas/libsatlas.so.3" - # Below libs are direct dependencies of libsatlas - LIBQUADMATH_PATH="/lib64/libquadmath.so.0" - fi MAYBE_LIB64=lib64 elif [[ "$OS_NAME" == *"Ubuntu"* ]]; then LIBGOMP_PATH="/usr/lib/x86_64-linux-gnu/libgomp.so.1" LIBNUMA_PATH="/usr/lib/x86_64-linux-gnu/libnuma.so.1" LIBELF_PATH="/usr/lib/x86_64-linux-gnu/libelf.so.1" - if [[ $ROCM_INT -ge 50300 ]]; then - LIBTINFO_PATH="/lib/x86_64-linux-gnu/libtinfo.so.6" - else - LIBTINFO_PATH="/lib/x86_64-linux-gnu/libtinfo.so.5" - fi + LIBTINFO_PATH="/lib/x86_64-linux-gnu/libtinfo.so.6" + LIBDW_PATH="/usr/lib/x86_64-linux-gnu/libdw.so.1" LIBDRM_PATH="/usr/lib/x86_64-linux-gnu/libdrm.so.2" LIBDRM_AMDGPU_PATH="/usr/lib/x86_64-linux-gnu/libdrm_amdgpu.so.1" - if [[ $ROCM_INT -ge 60100 && $ROCM_INT -lt 60300 ]]; then - # Below libs are direct dependencies of libhipsolver - LIBCHOLMOD_PATH="/lib/x86_64-linux-gnu/libcholmod.so.3" - # Below libs are direct dependencies of libcholmod - LIBSUITESPARSE_CONFIG_PATH="/lib/x86_64-linux-gnu/libsuitesparseconfig.so.5" - LIBAMD_PATH="/lib/x86_64-linux-gnu/libamd.so.2" - LIBCAMD_PATH="/lib/x86_64-linux-gnu/libcamd.so.2" - LIBCCOLAMD_PATH="/lib/x86_64-linux-gnu/libccolamd.so.2" - LIBCOLAMD_PATH="/lib/x86_64-linux-gnu/libcolamd.so.2" - LIBMETIS_PATH="/lib/x86_64-linux-gnu/libmetis.so.5" - LIBLAPACK_PATH="/lib/x86_64-linux-gnu/liblapack.so.3" - LIBBLAS_PATH="/lib/x86_64-linux-gnu/libblas.so.3" - # Below libs are direct dependencies of libblas - LIBGFORTRAN_PATH="/lib/x86_64-linux-gnu/libgfortran.so.5" - LIBQUADMATH_PATH="/lib/x86_64-linux-gnu/libquadmath.so.0" - fi MAYBE_LIB64=lib fi OS_SO_PATHS=($LIBGOMP_PATH $LIBNUMA_PATH\ $LIBELF_PATH $LIBTINFO_PATH\ + $LIBDW_PATH\ $LIBDRM_PATH $LIBDRM_AMDGPU_PATH\ $LIBSUITESPARSE_CONFIG_PATH\ $LIBCHOLMOD_PATH $LIBAMD_PATH\ diff --git a/.ci/onnx/README.md b/.ci/onnx/README.md deleted file mode 100644 index 47739136aabdf..0000000000000 --- a/.ci/onnx/README.md +++ /dev/null @@ -1,12 +0,0 @@ -# Jenkins - -The scripts in this directory are the entrypoint for testing ONNX exporter. - -The environment variable `BUILD_ENVIRONMENT` is expected to be set to -the build environment you intend to test. It is a hint for the build -and test scripts to configure Caffe2 a certain way and include/exclude -tests. Docker images, they equal the name of the image itself. For -example: `py2-cuda9.0-cudnn7-ubuntu16.04`. The Docker images that are -built on Jenkins and are used in triggered builds already have this -environment variable set in their manifest. Also see -`./docker/jenkins/*/Dockerfile` and search for `BUILD_ENVIRONMENT`. diff --git a/.ci/onnx/common.sh b/.ci/onnx/common.sh deleted file mode 100644 index b8f912fbbb4e6..0000000000000 --- a/.ci/onnx/common.sh +++ /dev/null @@ -1,107 +0,0 @@ -#!/bin/bash - -set -ex - -source "$(dirname "${BASH_SOURCE[0]}")/../pytorch/common_utils.sh" - -LOCAL_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) -ROOT_DIR=$(cd "$LOCAL_DIR"/../.. && pwd) -TEST_DIR="$ROOT_DIR/test" -pytest_reports_dir="${TEST_DIR}/test-reports/python" - -# Figure out which Python to use -PYTHON="$(which python)" -if [[ "${BUILD_ENVIRONMENT}" =~ py((2|3)\.?[0-9]?\.?[0-9]?) ]]; then - PYTHON=$(which "python${BASH_REMATCH[1]}") -fi - -if [[ "${BUILD_ENVIRONMENT}" == *rocm* ]]; then - # HIP_PLATFORM is auto-detected by hipcc; unset to avoid build errors - unset HIP_PLATFORM -fi - -mkdir -p "$pytest_reports_dir" || true - -########################################## -# copied from .ci/pytorch/common_utils.sh -########################################## - -function get_pinned_commit() { - cat .github/ci_commit_pins/"${1}".txt -} - -function pip_install_whl() { - # This is used to install PyTorch and other build artifacts wheel locally - # without using any network connection - - # Convert the input arguments into an array - local args=("$@") - - # Check if the first argument contains multiple paths separated by spaces - if [[ "${args[0]}" == *" "* ]]; then - # Split the string by spaces into an array - IFS=' ' read -r -a paths <<< "${args[0]}" - # Loop through each path and install individually - for path in "${paths[@]}"; do - echo "Installing $path" - python3 -mpip install --no-index --no-deps "$path" - done - else - # Loop through each argument and install individually - for path in "${args[@]}"; do - echo "Installing $path" - python3 -mpip install --no-index --no-deps "$path" - done - fi -} - -function pip_build_and_install() { - local build_target=$1 - local wheel_dir=$2 - - local found_whl=0 - for file in "${wheel_dir}"/*.whl - do - if [[ -f "${file}" ]]; then - found_whl=1 - break - fi - done - - # Build the wheel if it doesn't exist - if [ "${found_whl}" == "0" ]; then - python3 -m pip wheel \ - --no-build-isolation \ - --no-deps \ - -w "${wheel_dir}" \ - "${build_target}" - fi - - for file in "${wheel_dir}"/*.whl - do - pip_install_whl "${file}" - done -} - -function install_torchvision() { - local orig_preload - local commit - commit=$(get_pinned_commit vision) - orig_preload=${LD_PRELOAD} - if [ -n "${LD_PRELOAD}" ]; then - # Silence dlerror to work-around glibc ASAN bug, see https://sourceware.org/bugzilla/show_bug.cgi?id=27653#c9 - echo 'char* dlerror(void) { return "";}'|gcc -fpic -shared -o "${HOME}/dlerror.so" -x c - - LD_PRELOAD=${orig_preload}:${HOME}/dlerror.so - fi - - if [[ "${BUILD_ENVIRONMENT}" == *cuda* ]]; then - # Not sure if both are needed, but why not - export FORCE_CUDA=1 - export WITH_CUDA=1 - fi - pip_build_and_install "git+https://github.com/pytorch/vision.git@${commit}" dist/vision - - if [ -n "${LD_PRELOAD}" ]; then - LD_PRELOAD=${orig_preload} - fi -} diff --git a/.ci/onnx/test.sh b/.ci/onnx/test.sh deleted file mode 100755 index 1f2a23b49dc45..0000000000000 --- a/.ci/onnx/test.sh +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/bash - -# shellcheck source=./common.sh -source "$(dirname "${BASH_SOURCE[0]}")/common.sh" - -# Workaround for dind-rootless userid mapping (https://github.com/pytorch/ci-infra/issues/96) -WORKSPACE_ORIGINAL_OWNER_ID=$(stat -c '%u' "/var/lib/jenkins/workspace") -cleanup_workspace() { - echo "sudo may print the following warning message that can be ignored. The chown command will still run." - echo " sudo: setrlimit(RLIMIT_STACK): Operation not permitted" - echo "For more details refer to https://github.com/sudo-project/sudo/issues/42" - sudo chown -R "$WORKSPACE_ORIGINAL_OWNER_ID" /var/lib/jenkins/workspace -} -# Disable shellcheck SC2064 as we want to parse the original owner immediately. -# shellcheck disable=SC2064 -trap_add cleanup_workspace EXIT -sudo chown -R jenkins /var/lib/jenkins/workspace -git config --global --add safe.directory /var/lib/jenkins/workspace - -if [[ "$BUILD_ENVIRONMENT" == *onnx* ]]; then - # TODO: This can be removed later once vision is also part of the Docker image - install_torchvision - # JIT C++ extensions require ninja, so put it into PATH. - export PATH="/var/lib/jenkins/.local/bin:$PATH" - # NB: ONNX test is fast (~15m) so it's ok to retry it few more times to avoid any flaky issue, we - # need to bring this to the standard PyTorch run_test eventually. The issue will be tracked in - # https://github.com/pytorch/pytorch/issues/98626 - "$ROOT_DIR/scripts/onnx/test.sh" -fi diff --git a/.ci/pytorch/binary_populate_env.sh b/.ci/pytorch/binary_populate_env.sh index 9ea588d1555e2..53914914c8c93 100755 --- a/.ci/pytorch/binary_populate_env.sh +++ b/.ci/pytorch/binary_populate_env.sh @@ -112,9 +112,8 @@ if [[ "$PACKAGE_TYPE" =~ .*wheel.* && -n "$PYTORCH_BUILD_VERSION" && "$PYTORCH_B fi fi -USE_GLOO_WITH_OPENSSL="ON" +USE_GLOO_WITH_OPENSSL="OFF" if [[ "$GPU_ARCH_TYPE" =~ .*aarch64.* ]]; then - USE_GLOO_WITH_OPENSSL="OFF" USE_GOLD_LINKER="OFF" fi diff --git a/.ci/pytorch/build.sh b/.ci/pytorch/build.sh index 5853b7be2a98a..0914861b2d44c 100755 --- a/.ci/pytorch/build.sh +++ b/.ci/pytorch/build.sh @@ -144,6 +144,7 @@ if [[ "$BUILD_ENVIRONMENT" == *xpu* ]]; then export USE_XCCL=1 export USE_MPI=0 export TORCH_XPU_ARCH_LIST=pvc + export USE_STATIC_MKL=1 fi # sccache will fail for CUDA builds if all cores are used for compiling @@ -174,14 +175,9 @@ if [[ "$BUILD_ENVIRONMENT" == *cuda* ]] && echo "${TORCH_CUDA_ARCH_LIST}" | tr ' export BUILD_CUSTOM_STEP="ninja -C build flash_attention -j ${J}" fi -if [[ "${BUILD_ENVIRONMENT}" == *clang* ]]; then - export CC=clang - export CXX=clang++ - # TODO: Removeme once all the wrappers are gone - if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then - sudo rm -f /opt/cache/bin/clang++ - fi - +# TODO: Removeme once all the wrappers are gone +if [[ "$BUILD_ENVIRONMENT" == *clang* ]] && [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then + sudo rm -f /opt/cache/bin/clang++ fi if [[ "$BUILD_ENVIRONMENT" == *-clang*-asan* ]]; then diff --git a/.ci/pytorch/check_binary.sh b/.ci/pytorch/check_binary.sh index 3ff73084e4065..9356970394e81 100755 --- a/.ci/pytorch/check_binary.sh +++ b/.ci/pytorch/check_binary.sh @@ -292,25 +292,6 @@ if [[ "$PACKAGE_TYPE" != 'libtorch' ]]; then popd fi -############################################################################### -# Check PyTorch supports TCP_TLS gloo transport -############################################################################### - -if [[ "$(uname)" == 'Linux' && "$PACKAGE_TYPE" != 'libtorch' ]]; then - GLOO_CHECK="import torch.distributed as dist -try: - dist.init_process_group('gloo', rank=0, world_size=1) -except RuntimeError as e: - print(e) -" - RESULT=`GLOO_DEVICE_TRANSPORT=TCP_TLS MASTER_ADDR=localhost MASTER_PORT=63945 python -c "$GLOO_CHECK"` - GLOO_TRANSPORT_IS_NOT_SUPPORTED='gloo transport is not supported' - if [[ "$RESULT" =~ "$GLOO_TRANSPORT_IS_NOT_SUPPORTED" ]]; then - echo "PyTorch doesn't support TLS_TCP transport, please build with USE_GLOO_WITH_OPENSSL=1" - exit 1 - fi -fi - ############################################################################### # Restore LD_LIBRARY_PATH to its original value ############################################################################### diff --git a/.ci/pytorch/common.sh b/.ci/pytorch/common.sh index eae12816fe71e..94d9629eac519 100644 --- a/.ci/pytorch/common.sh +++ b/.ci/pytorch/common.sh @@ -5,7 +5,7 @@ source "$(dirname "${BASH_SOURCE[0]}")/common_utils.sh" set -ex -o pipefail -# for ROCm environment variables +# Source ROCm environment variables (paths may vary between tarball/wheel installs) if [[ "${BUILD_ENVIRONMENT}" == *rocm* ]] && [[ -f /etc/rocm_env.sh ]]; then # shellcheck disable=SC1091 source /etc/rocm_env.sh @@ -14,6 +14,19 @@ fi # Required environment variables: # $BUILD_ENVIRONMENT (should be set by your Docker image) +# Select compiler based on build environment name. Images that have both +# GCC and Clang installed default cc/c++ to Clang (via install_clang.sh), +# so we need to override when a gcc build is requested. +if [[ "${BUILD_ENVIRONMENT}" == *clang* ]]; then + export CC=clang + export CXX=clang++ +elif [[ "${BUILD_ENVIRONMENT}" == *gcc* ]]; then + export CC=gcc + export CXX=g++ + sudo update-alternatives --install /usr/bin/cc cc /usr/bin/gcc 100 + sudo update-alternatives --install /usr/bin/c++ c++ /usr/bin/g++ 100 +fi + # Figure out which Python to use for ROCm if [[ "${BUILD_ENVIRONMENT}" == *rocm* ]]; then # HIP_PLATFORM is auto-detected by hipcc; unset to avoid build errors diff --git a/.ci/pytorch/common_utils.sh b/.ci/pytorch/common_utils.sh index b5c0a5e43dea7..bb1eaeeb5e103 100644 --- a/.ci/pytorch/common_utils.sh +++ b/.ci/pytorch/common_utils.sh @@ -307,25 +307,9 @@ function install_torchao() { } function install_flash_attn_cute() { - echo "Installing FlashAttention CuTe from GitHub..." - # Grab latest main til we have a pinned commit - local flash_attn_commit - flash_attn_commit=$(git ls-remote https://github.com/Dao-AILab/flash-attention.git HEAD | cut -f1) - - # Clone the repo to a temporary directory - rm -rf flash-attention-build - git clone --depth 1 --recursive https://github.com/Dao-AILab/flash-attention.git flash-attention-build - - pushd flash-attention-build - git checkout "${flash_attn_commit}" - - # Install only the 'cute' sub-directory - pip_install flash_attn/cute/ - popd - - # remove the local repo - rm -rf flash-attention-build - echo "FlashAttention CuTe installation complete." + echo "Installing FlashAttention 4 from PyPI..." + pip_install flash-attn-4==4.0.0b5 + echo "FlashAttention 4 installation complete." } function install_cutlass_dsl() { diff --git a/.ci/pytorch/cpp_doc_push_script.sh b/.ci/pytorch/cpp_doc_push_script.sh index f085fa78bebe9..d0b4fd38826fa 100755 --- a/.ci/pytorch/cpp_doc_push_script.sh +++ b/.ci/pytorch/cpp_doc_push_script.sh @@ -1,7 +1,7 @@ #!/bin/bash # This is where the local pytorch install in the docker image is located -pt_checkout="/var/lib/jenkins/workspace" +pt_checkout="${GITHUB_WORKSPACE:-/var/lib/jenkins/workspace}" # Since we're cat-ing this file, we need to escape all $'s echo "cpp_doc_push_script.sh: Invoked with $*" @@ -60,6 +60,34 @@ time python tools/setup_helpers/generate_code.py \ pushd docs/cpp time make VERBOSE=1 html +# Run C++ API coverage check (allowlist-based + HTML formatting) +echo "Running C++ docs coverage check..." +python check_coverage.py --coverxygen || coverage_exit=$? + +# Generate coverxygen HTML report if coverxygen produced output +if [ -f coverxygen.info ] && command -v genhtml &> /dev/null; then + genhtml --no-function-coverage coverxygen.info -o build/html/_coverage \ + --title "PyTorch C++ API Doc Coverage" \ + --legend --highlight 2>/dev/null || true +fi + +# Copy coverage reports into the build output so they get uploaded +mkdir -p build/html/_coverage +cp -f cpp_coverage.txt cpp_html_issues.txt build/html/_coverage/ 2>/dev/null || true +cp -f coverxygen.info build/html/_coverage/ 2>/dev/null || true + +if [ "${coverage_exit:-0}" -ne 0 ]; then + echo "" + echo "========================================" + echo "C++ DOCS COVERAGE: HIGH-PRIORITY GAPS" + echo "========================================" + echo "" + cat cpp_coverage.txt + echo "" + echo "See the full coverage report at: _coverage/cpp_coverage.txt" + echo "See the HTML issues report at: _coverage/cpp_html_issues.txt" +fi + popd popd @@ -76,6 +104,7 @@ cp -r "${pt_checkout}"/docs/cpp/build/html/* . # Copy back _config.yml rm -rf _config.yml mv /tmp/cppdocs-sync/* . +touch .nojekyll # Make a new commit git add . || true diff --git a/.ci/pytorch/macos-build.sh b/.ci/pytorch/macos-build.sh index f2bcb486cf95b..3259c62149b61 100755 --- a/.ci/pytorch/macos-build.sh +++ b/.ci/pytorch/macos-build.sh @@ -48,7 +48,7 @@ if [[ ${BUILD_ENVIRONMENT} == *"distributed"* ]]; then else # Explicitly set USE_DISTRIBUTED=0 to align with the default build config on mac. This also serves as the sole CI config that tests # that building with USE_DISTRIBUTED=0 works at all. See https://github.com/pytorch/pytorch/issues/86448 - USE_DISTRIBUTED=0 USE_OPENMP=1 MACOSX_DEPLOYMENT_TARGET=11.0 WERROR=1 BUILD_TEST=OFF USE_PYTORCH_METAL=1 python -m build --wheel --no-isolation -C--build-option=--plat-name=macosx_11_0_arm64 + USE_DISTRIBUTED=0 USE_OPENMP=1 WERROR=1 BUILD_TEST=OFF USE_PYTORCH_METAL=1 python -m build --wheel --no-isolation fi if which sccache > /dev/null; then print_sccache_stats diff --git a/.ci/pytorch/macos-common.sh b/.ci/pytorch/macos-common.sh index 6826a52577a29..d21ee458f6c99 100755 --- a/.ci/pytorch/macos-common.sh +++ b/.ci/pytorch/macos-common.sh @@ -9,6 +9,6 @@ sysctl -a | grep machdep.cpu # These are required for both the build job and the test job. # In the latter to test cpp extensions. -export MACOSX_DEPLOYMENT_TARGET=11.1 +export MACOSX_DEPLOYMENT_TARGET=14.0 export CXX=clang++ export CC=clang diff --git a/.ci/pytorch/python_doc_push_script.sh b/.ci/pytorch/python_doc_push_script.sh index 6bcd46c4815a6..44b2a26e9fa66 100755 --- a/.ci/pytorch/python_doc_push_script.sh +++ b/.ci/pytorch/python_doc_push_script.sh @@ -1,7 +1,7 @@ #!/bin/bash # This is where the local pytorch install in the docker image is located -pt_checkout="/var/lib/jenkins/workspace" +pt_checkout="${GITHUB_WORKSPACE:-/var/lib/jenkins/workspace}" source "$pt_checkout/.ci/pytorch/common_utils.sh" @@ -50,8 +50,14 @@ echo "install_path: $install_path version: $version" build_docs () { set +e - set -o pipefail - make "$1" 2>&1 | tee /tmp/docs_build.txt + # Don't pipe through tee: sphinx -j auto forks workers that inherit + # the pipe fd and hold it open after sphinx exits, causing tee to + # block forever. Write to a file and tail with --pid so it exits + # (after draining) when make finishes. + make "$1" > /tmp/docs_build.txt 2>&1 & + local make_pid=$! + tail -f --pid=$make_pid /tmp/docs_build.txt + wait $make_pid code=$? if [ $code -ne 0 ]; then set +x @@ -87,7 +93,10 @@ pushd docs if [ "$is_main_doc" = true ]; then build_docs html || exit $? - make coverage + # Run coverage check without parallel workers since it's a quick + # check that doesn't need parallelism, and avoids re-triggering the + # expensive parallel read/write machinery. + SPHINXOPTS="-WT --keep-going" make coverage # Now we have the coverage report, we need to make sure it is empty. # Sphinx 7.2.6+ format: python.txt contains a statistics table with a TOTAL row # showing the undocumented count in the third column. diff --git a/.ci/pytorch/smoke_test/check_wheel_tags.py b/.ci/pytorch/smoke_test/check_wheel_tags.py index 47853183e0270..901304657b531 100644 --- a/.ci/pytorch/smoke_test/check_wheel_tags.py +++ b/.ci/pytorch/smoke_test/check_wheel_tags.py @@ -1,5 +1,4 @@ """Validate wheel platform tags and macOS dylib minos. - Supports two modes: 1. Pre-install: reads .whl files from PYTORCH_FINAL_PACKAGE_DIR 2. Post-install: reads metadata from installed torch package (soft warnings) @@ -11,6 +10,7 @@ import re import subprocess import sys +import tempfile import zipfile from pathlib import Path @@ -66,8 +66,17 @@ def check_wheel_platform_tag() -> None: if target_os == "linux" and platform.machine() == "aarch64": target_os = "linux-aarch64" expected_python = f"cp{sys.version_info.major}{sys.version_info.minor}" + import sysconfig + abiflags = getattr(sys, "abiflags", "") + if not abiflags and ( + os.getenv("MATRIX_PYTHON_VERSION", "").endswith("t") + or bool(sysconfig.get_config_var("Py_GIL_DISABLED")) + or not getattr(sys, "_is_gil_enabled", lambda: True)() + ): + abiflags = "t" expected_abi = f"cp{sys.version_info.major}{sys.version_info.minor}{abiflags}" + print(f"Expected ABI tag: {expected_abi}") platform_pattern = EXPECTED_PLATFORM_TAGS.get(target_os) if not platform_pattern: @@ -113,7 +122,6 @@ def check_wheel_platform_tag() -> None: f"expected format: --" ) raise RuntimeError(msg) - continue python_tag, abi_tag, platform_tag = parts @@ -143,92 +151,117 @@ def check_wheel_platform_tag() -> None: print(f"OK: Wheel tag(s) valid for {source}: {', '.join(tags)}") -def check_mac_wheel_minos() -> None: - """Check that dylib minos matches the wheel platform tag on macOS. +def _check_dylibs_minos(dylibs: list, expected_minos: str, source: str) -> None: + mismatches = [] + for dylib in dylibs: + try: + result = subprocess.run( + ["otool", "-l", str(dylib)], + capture_output=True, + text=True, + timeout=30, + ) + except Exception: + continue - Extracts dylibs from the .whl in PYTORCH_FINAL_PACKAGE_DIR to a temp dir, - then verifies each dylib's minos (from otool -l) matches the platform tag. - """ + minos = None + lines = result.stdout.splitlines() + for i, line in enumerate(lines): + s = line.strip() + if "LC_BUILD_VERSION" in s: + for j in range(i + 1, min(i + 6, len(lines))): + if lines[j].strip().startswith("minos"): + minos = lines[j].strip().split()[1] + break + break + if "LC_VERSION_MIN_MACOSX" in s: + for j in range(i + 1, min(i + 4, len(lines))): + if lines[j].strip().startswith("version"): + minos = lines[j].strip().split()[1] + break + break + + # A dylib with a lower minos than the wheel tag is safe (forward compatible). + # Only flag dylibs that require a *higher* macOS than the wheel claims to support. + if minos and tuple(int(x) for x in minos.split(".")) > tuple( + int(x) for x in expected_minos.split(".") + ): + mismatches.append( + f"{dylib.name}: minos={minos}, expected<={expected_minos}" + ) + + if mismatches: + raise RuntimeError( + f"minos/platform tag mismatch in {len(mismatches)} dylib(s):\n" + + "\n".join(f" {m}" for m in mismatches) + ) + print( + f"OK: All {len(dylibs)} dylib(s) have minos matching " + f"platform tag ({expected_minos}) for {source}" + ) + + +def check_mac_wheel_minos() -> None: if sys.platform != "darwin": return wheel_dir = os.getenv("PYTORCH_FINAL_PACKAGE_DIR", "") - if not wheel_dir or not os.path.isdir(wheel_dir): - print("PYTORCH_FINAL_PACKAGE_DIR not set, skipping wheel minos check") - return - whls = list(Path(wheel_dir).glob("*.whl")) - if not whls: - print(f"No .whl files in {wheel_dir}, skipping wheel minos check") - return + if wheel_dir and os.path.isdir(wheel_dir): + # Mode 1: extract dylibs from .whl file + whls = list(Path(wheel_dir).glob("*.whl")) + if not whls: + print(f"No .whl files in {wheel_dir}, skipping wheel minos check") + return + + macos_whl_re = re.compile(r"macosx_(\d+)_(\d+)_(\w+)\.whl$") + for whl in whls: + print(f"Checking wheel tag minos for: {whl.name}") + m = macos_whl_re.search(whl.name) + if not m: + print(f"No macOS platform tag in {whl.name}, skipping") + continue + expected_minos = f"{m.group(1)}.{m.group(2)}" + + with tempfile.TemporaryDirectory() as tmpdir: + with zipfile.ZipFile(whl, "r") as zf: + dylib_names = [n for n in zf.namelist() if n.endswith(".dylib")] + if not dylib_names: + print("No .dylib files in wheel, skipping minos check") + continue + for name in dylib_names: + zf.extract(name, tmpdir) + dylibs = list(Path(tmpdir).rglob("*.dylib")) + _check_dylibs_minos(dylibs, expected_minos, whl.name) + else: + # Mode 2: read from installed torch package + print("PYTORCH_FINAL_PACKAGE_DIR not set, checking installed torch dylibs") + try: + tags = _extract_installed_wheel_tags("torch") + except Exception as e: + print(f"Could not read installed torch metadata: {e}, skipping") + return - import tempfile + expected_minos = None + for tag_str in tags: + m = re.search(r"macosx_(\d+)_(\d+)_\w+", tag_str) + if m: + expected_minos = f"{m.group(1)}.{m.group(2)}" + break - for whl in whls: - print(f"Checking wheel tag minos for: {whl.name}") + if not expected_minos: + print("No macOS platform tag found in installed torch metadata, skipping") + return - m = re.search(r"macosx_(\d+)_(\d+)_(\w+)\.whl$", whl.name) - if not m: - print(f"No macOS platform tag in {whl.name}, skipping") - continue + print(f"Expected minos from installed wheel tag: {expected_minos}") - expected_minos = f"{m.group(1)}.{m.group(2)}" - print(f"Expected minos from platform tag: {expected_minos}") - - # Extract dylibs from wheel to temp dir - with tempfile.TemporaryDirectory() as tmpdir: - with zipfile.ZipFile(whl, "r") as zf: - dylib_names = [n for n in zf.namelist() if n.endswith(".dylib")] - if not dylib_names: - print("No .dylib files in wheel, skipping minos check") - continue - for name in dylib_names: - zf.extract(name, tmpdir) - - dylibs = list(Path(tmpdir).rglob("*.dylib")) - mismatches = [] - for dylib in dylibs: - try: - result = subprocess.run( - ["otool", "-l", str(dylib)], - capture_output=True, - text=True, - timeout=30, - ) - except Exception: - continue - - minos = None - lines = result.stdout.splitlines() - for i, line in enumerate(lines): - s = line.strip() - if "LC_BUILD_VERSION" in s: - for j in range(i + 1, min(i + 6, len(lines))): - if lines[j].strip().startswith("minos"): - minos = lines[j].strip().split()[1] - break - break - if "LC_VERSION_MIN_MACOSX" in s: - for j in range(i + 1, min(i + 4, len(lines))): - if lines[j].strip().startswith("version"): - minos = lines[j].strip().split()[1] - break - break + import torch - if minos and minos != expected_minos: - mismatches.append( - f"{dylib.name}: minos={minos}, expected={expected_minos}" - ) - - if mismatches: - raise RuntimeError( - f"minos/platform tag mismatch in {len(mismatches)} dylib(s):\n" - + "\n".join(f" {m}" for m in mismatches) - ) - print( - f"OK: All {len(dylibs)} dylib(s) have minos matching " - f"platform tag ({expected_minos})" - ) + torch_dir = Path(torch.__file__).parent + dylibs = list(torch_dir.rglob("*.dylib")) + if not dylibs: + raise RuntimeError("No .dylib files found in installed torch") + _check_dylibs_minos(dylibs, expected_minos, "installed torch") if __name__ == "__main__": diff --git a/.ci/pytorch/test.sh b/.ci/pytorch/test.sh index 772870dc5ba59..0e1087cf372ef 100755 --- a/.ci/pytorch/test.sh +++ b/.ci/pytorch/test.sh @@ -45,6 +45,16 @@ if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then fi fi +# Remove onnxruntime if present to avoid interference with non-ONNX tests +if [[ "$TEST_CONFIG" != "onnx" ]]; then + pip uninstall -y onnxruntime 2>/dev/null || true +fi + +# Remove dill to test that serialization works without it +if [[ "$BUILD_ENVIRONMENT" == *py3.10-gcc11 ]]; then + pip uninstall -y dill 2>/dev/null || true +fi + echo "Environment variables:" env @@ -145,11 +155,34 @@ env echo "Testing pytorch" +# Set OMP_NUM_THREADS to nproc/4 on k8s ARC runners if not already set. +# +# We use nproc (cgroup-aware) rather than os.cpu_count() because on k8s (ARC) +# pods, os.cpu_count() returns the host's CPU count (e.g., 192) rather than +# the pod's cpuset allocation (e.g., 16). +# +# We use nproc/4 rather than nproc because OpenMP spin-waits at thread barriers. +# When thread count equals cpuset size (e.g., 16 threads on 16 CPUs), spinning +# barrier threads monopolize all CPUs and the OS must context-switch to let +# actual work complete. This causes ~5000x slowdowns on small tensor ops +# (e.g., aten::copy_ on 147KB: ~34ms instead of ~7us). Using nproc/4 leaves +# headroom for the main thread and for NUM_PROCS=3 parallel test processes. +if [[ -z "${OMP_NUM_THREADS:-}" ]] && [[ -n "${USE_ARC:-}" ]]; then + OMP_NUM_THREADS=$(( $(nproc) / 4 )) + # Floor of 4: low OMP_NUM_THREADS (1-2) changes floating-point reduction + # order, causing numerical mismatches in tests with tight tolerances + # (e.g., test_batchnorm_nhwc_cpu). + if [[ "$OMP_NUM_THREADS" -lt 4 ]]; then + OMP_NUM_THREADS=4 + fi + export OMP_NUM_THREADS +fi + export LANG=C.UTF-8 PR_NUMBER=${PR_NUMBER:-${CIRCLE_PR_NUMBER:-}} -if [[ -d "${HF_CACHE}" ]]; then +if [[ -d "${HF_CACHE}" && "$TEST_CONFIG" != "onnx" ]]; then export HF_HOME="${HF_CACHE}" fi @@ -367,9 +400,11 @@ test_python_smoke_b200() { assert_git_not_dirty } + test_python_smoke_xpu() { # Smoke tests for XPU client time python test/run_test.py --include test_transformers $PYTHON_TEST_EXTRA_OPTION --upload-artifacts-while-running + time test_xpu_sycl_tla_backend assert_git_not_dirty } @@ -421,6 +456,15 @@ test_h100_cutlass_backend() { TORCHINDUCTOR_CUTLASS_DIR=$(realpath "./third_party/cutlass") python test/run_test.py --include inductor/test_cutlass_evt $PYTHON_TEST_EXTRA_OPTION --upload-artifacts-while-running } +test_xpu_sycl_tla_backend() { + # Inductor sycl-tla backend tests for XPU + # shellcheck disable=SC1091 + source /opt/intel/oneapi/mkl/latest/env/vars.sh + sycl_tla_dir=$(realpath "./third_party/sycl-tla") + rm -rf "${sycl_tla_dir}" && git clone --depth 1 --single-branch -b v0.8 --quiet https://github.com/intel/sycl-tla.git "${sycl_tla_dir}" + TORCHINDUCTOR_CUTLASS_DIR=$(realpath "./third_party/sycl-tla") python test/run_test.py --include inductor/test_cutlass_backend -k "not addmm" $PYTHON_TEST_EXTRA_OPTION --upload-artifacts-while-running +} + test_lazy_tensor_meta_reference_disabled() { export TORCH_DISABLE_FUNCTIONALIZATION_META_REFERENCE=1 echo "Testing lazy tensor operations without meta reference" @@ -437,12 +481,16 @@ test_dynamo_core() { } test_dynamo_cpython() { + # Disable TD for cpython since it's pretty cheap to run the cpython tests (< 10 min) + # and if TD is enabled, only 25% of the tests will be executed + export NO_TD=1 time python test/run_test.py \ --include-cpython-tests \ --dynamo \ --verbose \ --upload-artifacts-while-running assert_git_not_dirty + unset NO_TD } test_dynamo_wrapped_shard() { @@ -549,7 +597,7 @@ test_inductor_shard() { # Do not add --inductor for the following inductor unit tests, otherwise we will fail because of nested dynamo state python test/run_test.py \ - --include inductor/test_torchinductor inductor/test_torchinductor_opinfo inductor/test_aot_inductor \ + --include inductor/test_torchinductor inductor/test_torchinductor_opinfo inductor/test_aot_inductor inductor/test_cpu_select_algorithm \ --shard "$1" "$NUM_TEST_SHARDS" \ --verbose } @@ -615,8 +663,10 @@ test_inductor_cpp_wrapper_shard() { -k 'take' \ --shard "$1" "$NUM_TEST_SHARDS" \ --verbose - TORCHINDUCTOR_AUTOTUNE_AT_COMPILE_TIME=0 python test/run_test.py \ - --include inductor/test_torchinductor inductor/test_triton_kernels\ + # Keep testing TORCHINDUCTOR_AUTOTUNE_AT_COMPILE_TIME=1 for the near future. + # Will drop this after AOTInductor also switches to lazy Triton compilation. + TORCHINDUCTOR_AUTOTUNE_AT_COMPILE_TIME=1 python test/run_test.py \ + --include inductor/test_torchinductor inductor/test_triton_kernels inductor/test_max_autotune \ --shard "$1" "$NUM_TEST_SHARDS" \ --verbose if [[ "${BUILD_ENVIRONMENT}" == *xpu* ]]; then @@ -732,6 +782,10 @@ test_perf_for_dashboard() { TEST_REPORTS_DIR=$(pwd)/test/test-reports mkdir -p "$TEST_REPORTS_DIR" + if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" ]]; then + mkdir -p "$TEST_REPORTS_DIR/profiler_traces" + fi + local suite="$1" shift @@ -791,52 +845,97 @@ test_perf_for_dashboard() { fi if [[ "$DASHBOARD_TAG" == *default-true* ]]; then + local profiler_trace_flags=() + if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then + profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_no_cudagraphs_${suite}_${dtype}_${mode}_${device}") + fi $TASKSET python "benchmarks/dynamo/$suite.py" \ "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" --disable-cudagraphs "$@" \ + "${profiler_trace_flags[@]}" \ --output "$TEST_REPORTS_DIR/${backend}_no_cudagraphs_${suite}_${dtype}_${mode}_${device}_${target}.csv" fi if [[ "$DASHBOARD_TAG" == *cudagraphs-true* ]]; then + local profiler_trace_flags=() + if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then + profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_with_cudagraphs_${suite}_${dtype}_${mode}_${device}") + fi $TASKSET python "benchmarks/dynamo/$suite.py" \ "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" "$@" \ + "${profiler_trace_flags[@]}" \ --output "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_${suite}_${dtype}_${mode}_${device}_${target}.csv" fi if [[ "$DASHBOARD_TAG" == *dynamic-true* ]]; then + local profiler_trace_flags=() + if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then + profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_dynamic_${suite}_${dtype}_${mode}_${device}") + fi $TASKSET python "benchmarks/dynamo/$suite.py" \ "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" --dynamic-shapes \ --dynamic-batch-only "$@" \ + "${profiler_trace_flags[@]}" \ --output "$TEST_REPORTS_DIR/${backend}_dynamic_${suite}_${dtype}_${mode}_${device}_${target}.csv" fi if [[ "$DASHBOARD_TAG" == *cppwrapper-true* ]]; then + local profiler_trace_flags=() + if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then + profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_cpp_wrapper_${suite}_${dtype}_${mode}_${device}") + fi TORCHINDUCTOR_CPP_WRAPPER=1 $TASKSET python "benchmarks/dynamo/$suite.py" \ "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" --disable-cudagraphs "$@" \ + "${profiler_trace_flags[@]}" \ --output "$TEST_REPORTS_DIR/${backend}_cpp_wrapper_${suite}_${dtype}_${mode}_${device}_${target}.csv" fi if [[ "$DASHBOARD_TAG" == *freezing_cudagraphs-true* ]] && [[ "$mode" == "inference" ]]; then + local profiler_trace_flags=() + if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then + profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_with_cudagraphs_freezing_${suite}_${dtype}_${mode}_${device}") + fi $TASKSET python "benchmarks/dynamo/$suite.py" \ "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" "$@" --freezing \ + "${profiler_trace_flags[@]}" \ --output "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_freezing_${suite}_${dtype}_${mode}_${device}_${target}.csv" fi if [[ "$DASHBOARD_TAG" == *freeze_autotune_cudagraphs-true* ]] && [[ "$mode" == "inference" ]]; then + local profiler_trace_flags=() + if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then + profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_with_cudagraphs_freezing_autotune_${suite}_${dtype}_${mode}_${device}") + fi TORCHINDUCTOR_MAX_AUTOTUNE=1 $TASKSET python "benchmarks/dynamo/$suite.py" \ "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" "$@" --freezing \ + "${profiler_trace_flags[@]}" \ --output "$TEST_REPORTS_DIR/${backend}_with_cudagraphs_freezing_autotune_${suite}_${dtype}_${mode}_${device}_${target}.csv" fi if [[ "$DASHBOARD_TAG" == *aotinductor-true* ]] && [[ "$mode" == "inference" ]]; then - if [[ "$target" == "accuracy" ]]; then # Also collect Export pass rate and display as a separate row + if [[ "$target" == "accuracy" ]]; then $TASKSET python "benchmarks/dynamo/$suite.py" \ "${target_flag[@]}" --"$mode" --"$dtype" --export --disable-cudagraphs "$@" \ --output "$TEST_REPORTS_DIR/${backend}_export_${suite}_${dtype}_${mode}_${device}_${target}.csv" fi + local profiler_trace_flags=() + if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then + profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_aot_inductor_${suite}_${dtype}_${mode}_${device}") + fi $TASKSET python "benchmarks/dynamo/$suite.py" \ "${target_flag[@]}" --"$mode" --"$dtype" --export-aot-inductor --disable-cudagraphs "$@" \ + "${profiler_trace_flags[@]}" \ --output "$TEST_REPORTS_DIR/${backend}_aot_inductor_${suite}_${dtype}_${mode}_${device}_${target}.csv" fi if [[ "$DASHBOARD_TAG" == *maxautotune-true* ]]; then + local profiler_trace_flags=() + if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" && "$target" == "performance" ]]; then + profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${backend}_max_autotune_${suite}_${dtype}_${mode}_${device}") + fi TORCHINDUCTOR_MAX_AUTOTUNE=1 $TASKSET python "benchmarks/dynamo/$suite.py" \ "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" "$@" \ + "${profiler_trace_flags[@]}" \ --output "$TEST_REPORTS_DIR/${backend}_max_autotune_${suite}_${dtype}_${mode}_${device}_${target}.csv" fi + if [[ "$DASHBOARD_TAG" == *deterministic_perf-true* ]]; then + $TASKSET python "benchmarks/dynamo/$suite.py" \ + "${target_flag[@]}" --"$mode" --"$dtype" --backend "$backend" --disable-cudagraphs --deterministic "$@" \ + --output "$TEST_REPORTS_DIR/${backend}_deterministic_perf_${suite}_${dtype}_${mode}_${device}_${target}.csv" + fi done done } @@ -864,9 +963,15 @@ test_single_dynamo_benchmark() { fi if [[ "${TEST_CONFIG}" == *perf_compare* ]]; then + local profiler_trace_flags=() + if [[ "${EXPORT_PROFILER_TRACE:-0}" == "1" ]]; then + mkdir -p "$TEST_REPORTS_DIR/profiler_traces" + profiler_trace_flags=(--export-profiler-trace --profiler-trace-name "$TEST_REPORTS_DIR/profiler_traces/${name}_${suite}") + fi python "benchmarks/dynamo/$suite.py" \ --ci --performance --disable-cudagraphs --inductor \ "${DYNAMO_BENCHMARK_FLAGS[@]}" "$@" "${partition_flags[@]}" \ + "${profiler_trace_flags[@]}" \ --output "$TEST_REPORTS_DIR/${name}_${suite}.csv" elif [[ "${TEST_CONFIG}" == *perf* ]]; then test_perf_for_dashboard "$suite" \ @@ -920,6 +1025,49 @@ test_inductor_triton_cpu() { assert_git_not_dirty } +setup_torch_trace() { + if [[ "${ENABLE_TORCH_TRACE:-0}" != "1" ]]; then + return + fi + local trace_dir="${RUNNER_TEMP:-/tmp}/torch_traces" + mkdir -p "$trace_dir" + export TORCH_TRACE="$trace_dir" + echo "TORCH_TRACE enabled: writing structured trace logs to $trace_dir" +} + +collect_tlparse_output() { + if [[ "${ENABLE_TORCH_TRACE:-0}" != "1" ]]; then + return + fi + local trace_dir="${RUNNER_TEMP:-/tmp}/torch_traces" + local test_reports_dir + test_reports_dir=$(pwd)/test/test-reports + + if [[ ! -d "$trace_dir" ]] || [[ -z "$(ls -A "$trace_dir" 2>/dev/null)" ]]; then + echo "No torch trace files found in $trace_dir, skipping tlparse" + return + fi + + echo "Collecting tlparse output from $trace_dir" + + # Install tlparse if not already available + if ! command -v tlparse &>/dev/null; then + pip install tlparse 2>/dev/null || { + echo "Warning: failed to install tlparse, skipping HTML generation" + return + } + fi + + # Run tlparse to generate HTML report + mkdir -p "$test_reports_dir/tlparse_output" + tlparse -o "$test_reports_dir/tlparse_output/" --no-browser --overwrite "$trace_dir" 2>&1 || { + echo "Warning: tlparse failed to generate HTML output" + return + } + + echo "TLParse output generated in $test_reports_dir/tlparse_output/" +} + test_dynamo_benchmark() { # Usage: test_dynamo_benchmark huggingface 0 TEST_REPORTS_DIR=$(pwd)/test/test-reports @@ -1002,6 +1150,147 @@ test_inductor_torchbench_smoketest_perf() { done } +test_unbacked_parity_smoketest() { + # Check that unbacked batch-only has performance parity with backed batch-only + # Fails if any model regresses >THRESHOLD% consistently across 3 retries + TEST_REPORTS_DIR=$(pwd)/test/test-reports + mkdir -p "$TEST_REPORTS_DIR" + + local THRESHOLD=1.0 + local MAX_RETRIES=3 + local MODELS="MobileBertForMaskedLM|DistilBertForMaskedLM|DistillGPT2|T5Small" + + # Issue 6: Write per-run output files for post-failure debugging + run_comparison() { + local run_num=$1 + local output_file="$TEST_REPORTS_DIR/unbacked_parity_results_run${run_num}.txt" + python benchmarks/dynamo/huggingface.py \ + --compare-backed-unbacked \ + --performance --inference --inductor --device cuda \ + --filter "$MODELS" 2>&1 | tee "$output_file" + } + + check_regressions() { + local run_num=$1 + local output_file="$TEST_REPORTS_DIR/unbacked_parity_results_run${run_num}.txt" + # Parse the comparison table and check for regressions > threshold + # Returns 0 if regressions found, 1 if no regressions + local regressions=() + while IFS= read -r line; do + # Issue 3: Broadened regex to match model names with hyphens, slashes, dots + # Match lines like: " ModelName 10.000 10.500 +5.0%" + if [[ "$line" =~ ^[[:space:]]+([A-Za-z0-9_./-]+)[[:space:]]+([0-9.]+)[[:space:]]+([0-9.]+)[[:space:]]+\+([0-9.]+)% ]]; then + local model="${BASH_REMATCH[1]}" + local diff="${BASH_REMATCH[4]}" + # Nit: Use awk instead of bc -l to avoid dependency on bc + if awk "BEGIN{exit !($diff > $THRESHOLD)}"; then + regressions+=("$model:+${diff}%") + fi + fi + done < "$output_file" + + if [[ ${#regressions[@]} -gt 0 ]]; then + echo "Regressions found: ${regressions[*]}" + return 0 + fi + return 1 + } + + check_failures() { + local run_num=$1 + local output_file="$TEST_REPORTS_DIR/unbacked_parity_results_run${run_num}.txt" + # Issue 2: Check for any model failure — not just paired failures. + # Specifically flags when unbacked fails but backed succeeds (regression signal). + # Returns 0 if failures found, 1 if no failures + local current_model="" + local backed_failed=false + local unbacked_failed=false + local both_failures=() + local unbacked_only_failures=() + + # Append a sentinel header so the loop naturally evaluates the last real model + while IFS= read -r line; do + if [[ "$line" =~ ^---[[:space:]]+([A-Za-z0-9_./-]+)[[:space:]]+--- ]]; then + if [[ -n "$current_model" ]]; then + if $backed_failed && $unbacked_failed; then + both_failures+=("$current_model") + elif $unbacked_failed && ! $backed_failed; then + unbacked_only_failures+=("$current_model") + fi + fi + current_model="${BASH_REMATCH[1]}" + backed_failed=false + unbacked_failed=false + elif [[ "$line" =~ backed.*FAILED|backed.*TIMEOUT|backed.*ERROR ]]; then + backed_failed=true + elif [[ "$line" =~ unbacked.*FAILED|unbacked.*TIMEOUT|unbacked.*ERROR ]]; then + unbacked_failed=true + fi + done < <(cat "$output_file"; echo "--- END ---") + + local has_failures=false + if [[ ${#both_failures[@]} -gt 0 ]]; then + echo "❌ FAILURES DETECTED: Both backed and unbacked failed for: ${both_failures[*]}" + has_failures=true + fi + if [[ ${#unbacked_only_failures[@]} -gt 0 ]]; then + echo "❌ FAILURES DETECTED: Unbacked failed (but backed succeeded) for: ${unbacked_only_failures[*]}" + has_failures=true + fi + + if $has_failures; then + return 0 + fi + return 1 + } + + # Run initial comparison + echo "=== Run 1/$MAX_RETRIES ===" + run_comparison 1 + + # Check for failures first + if check_failures 1; then + echo "❌ Test failed: Models failed to run (see above for details)" + exit 1 + fi + + # Check for regressions + if ! check_regressions 1; then + echo "✅ PASSED: No regressions above ${THRESHOLD}% threshold" + exit 0 + fi + + # Regression detected - retry to confirm + local regression_count=1 + for ((retry=2; retry<=MAX_RETRIES; retry++)); do + echo "" + echo "=== Retry $retry/$MAX_RETRIES (potential regression detected) ===" + run_comparison "$retry" + + # Issue 4: Also check for failures on retries (e.g., intermittent OOM) + if check_failures "$retry"; then + echo "❌ Test failed: Models failed on retry $retry (see above for details)" + exit 1 + fi + + if check_regressions "$retry"; then + ((regression_count++)) + fi + done + + # Check if regression was consistent (majority of runs) + local required=$((MAX_RETRIES / 2 + 1)) + if [[ $regression_count -ge $required ]]; then + echo "" + echo "❌ REGRESSION CONFIRMED: Detected in $regression_count/$MAX_RETRIES runs (threshold: ${THRESHOLD}%)" + exit 1 + else + echo "" + echo "✅ PASSED: Regressions were not consistent ($regression_count/$MAX_RETRIES runs, needed $required)" + exit 0 + fi +} + test_inductor_set_cpu_affinity(){ JEMALLOC_LIB="$(find /usr/lib -name libjemalloc.so.2)" export LD_PRELOAD="$JEMALLOC_LIB":"$LD_PRELOAD" @@ -1513,7 +1802,9 @@ EOF fi # Ensure invalid item is in the test output. - echo "${test_output}" | grep -q "${invalid_item_name}" && ret=$? || ret=$? + # Use a here-string instead of a pipe to avoid SIGPIPE when grep -q + # exits early on large output (causes exit code 141 with pipefail). + grep -q "${invalid_item_name}" <<< "${test_output}" && ret=$? || ret=$? if [ $ret -ne 0 ]; then cat << EOF @@ -1673,7 +1964,10 @@ test_bazel() { tools/bazel test --config=cpu-only --test_timeout=480 --test_output=all --test_tag_filters=-gpu-required --test_filter=-*CUDA \ --no//c10:use_gflags --no//c10:use_glog //c10/... - tools/bazel test --config=cpu-only --test_timeout=480 --test_output=all --test_tag_filters=-gpu-required --test_filter=-*CUDA :all_tests + # rnn_test uses a convergence-based training loop (test_RNN_xor) whose + # runtime varies across CPU microarchitectures. On OSDC AVX-512 runners + # it can exceed 480s, so give extra headroom. + tools/bazel test --config=cpu-only --test_timeout=900 --test_output=all --test_tag_filters=-gpu-required --test_filter=-*CUDA :all_tests else # Increase the test timeout to 480 like CPU tests because modules_test frequently timeout tools/bazel test --test_timeout=480 --test_output=errors \ @@ -1795,37 +2089,6 @@ test_executorch() { assert_git_not_dirty } -test_linux_aarch64() { - python test/run_test.py --include test_modules test_utils test_mkldnn test_mkldnn_fusion test_openmp test_torch test_dynamic_shapes \ - test_transformers test_multiprocessing test_numpy_interop test_autograd test_binary_ufuncs test_complex test_spectral_ops \ - test_foreach test_reductions test_unary_ufuncs test_tensor_creation_ops test_ops profiler/test_memory_profiler \ - distributed/elastic/timer/api_test distributed/elastic/timer/local_timer_example distributed/elastic/timer/local_timer_test \ - test_linalg \ - test_jit test_jit_autocast test_ops_jit \ - test_nn nn/test_convolution functorch/test_ops functorch/test_aotdispatch \ - --shard "$SHARD_NUMBER" "$NUM_TEST_SHARDS" --verbose - - # Dynamo tests - python test/run_test.py --include dynamo/test_compile dynamo/test_backends dynamo/test_comptime dynamo/test_config \ - dynamo/test_functions dynamo/test_fx_passes_pre_grad dynamo/test_interop dynamo/test_model_output dynamo/test_modules \ - dynamo/test_optimizers dynamo/test_recompile_ux dynamo/test_recompiles \ - --shard "$SHARD_NUMBER" "$NUM_TEST_SHARDS" --verbose - - # Inductor tests - python test/run_test.py --include inductor/test_torchinductor inductor/test_benchmark_fusion inductor/test_codecache \ - inductor/test_config inductor/test_control_flow inductor/test_coordinate_descent_tuner inductor/test_fx_fusion \ - inductor/test_group_batch_fusion inductor/test_inductor_freezing inductor/test_inductor_utils \ - inductor/test_inplacing_pass inductor/test_kernel_benchmark inductor/test_layout_optim \ - inductor/test_max_autotune inductor/test_memory_planning inductor/test_metrics inductor/test_multi_kernel inductor/test_pad_mm \ - inductor/test_pattern_matcher inductor/test_perf inductor/test_profiler inductor/test_select_algorithm inductor/test_smoke \ - inductor/test_split_cat_fx_passes inductor/test_compile inductor/test_torchinductor \ - inductor/test_torchinductor_codegen_dynamic_shapes inductor/test_torchinductor_dynamic_shapes inductor/test_memory \ - inductor/test_triton_cpu_backend inductor/test_triton_extension_backend inductor/test_mkldnn_pattern_matcher inductor/test_cpu_cpp_wrapper \ - inductor/test_cpu_select_algorithm inductor/test_cpu_repro \ - inductor/test_aot_inductor inductor/test_fused_attention \ - --shard "$SHARD_NUMBER" "$NUM_TEST_SHARDS" --verbose -} - test_operator_benchmark() { TEST_REPORTS_DIR=$(pwd)/test/test-reports mkdir -p "$TEST_REPORTS_DIR" @@ -1897,7 +2160,10 @@ if ! [[ "${BUILD_ENVIRONMENT}" == *libtorch* || "${BUILD_ENVIRONMENT}" == *-baze (cd test && python -c "import torch; print(torch.__config__.show())") (cd test && python -c "import torch; print(torch.__config__.parallel_info())") fi -if [[ "${TEST_CONFIG}" == *numpy_2* ]]; then +if [[ "${TEST_CONFIG}" == "onnx" ]]; then + install_torchvision + "$(dirname "${BASH_SOURCE[0]}")/../../scripts/onnx/test.sh" +elif [[ "${TEST_CONFIG}" == *numpy_2* ]]; then # Install numpy-2.0.2 and compatible scipy & numba versions # Force re-install of pandas to avoid error where pandas checks numpy version from initial install and fails upon import TMP_PANDAS_VERSION=$(python -c "import pandas; print(pandas.__version__)" 2>/dev/null) @@ -1907,8 +2173,6 @@ if [[ "${TEST_CONFIG}" == *numpy_2* ]]; then python -m pip install --pre numpy==2.0.2 scipy==1.13.1 numba==0.60.0 fi python test/run_test.py --include dynamo/test_functions.py dynamo/test_unspec.py test_binary_ufuncs.py test_fake_tensor.py test_linalg.py test_numpy_interop.py test_tensor_creation_ops.py test_torch.py torch_np/test_basic.py -elif [[ "${BUILD_ENVIRONMENT}" == *aarch64* && "${TEST_CONFIG}" == 'default' ]]; then - test_linux_aarch64 elif [[ "${TEST_CONFIG}" == *backward* ]]; then test_forward_backward_compatibility # Do NOT add tests after bc check tests, see its comment. @@ -1959,7 +2223,9 @@ elif [[ "${TEST_CONFIG}" == *operator_microbenchmark* ]]; then elif [[ "${TEST_CONFIG}" == *attention_microbenchmark* ]]; then test_attention_microbenchmark elif [[ "${TEST_CONFIG}" == *inductor_distributed* ]]; then + setup_torch_trace test_inductor_distributed + collect_tlparse_output elif [[ "${TEST_CONFIG}" == *inductor-halide* ]]; then test_inductor_halide elif [[ "${TEST_CONFIG}" == *inductor-pallas* ]]; then @@ -1973,11 +2239,19 @@ elif [[ "${TEST_CONFIG}" == *aoti_cross_compile_for_windows* ]]; then elif [[ "${TEST_CONFIG}" == *huggingface* ]]; then install_torchvision id=$((SHARD_NUMBER-1)) - test_dynamo_benchmark huggingface "$id" + setup_torch_trace + if [[ "${TEST_CONFIG}" == *unbacked_parity* ]]; then + test_unbacked_parity_smoketest + else + test_dynamo_benchmark huggingface "$id" + fi + collect_tlparse_output elif [[ "${TEST_CONFIG}" == *timm* ]]; then install_torchvision id=$((SHARD_NUMBER-1)) + setup_torch_trace test_dynamo_benchmark timm_models "$id" + collect_tlparse_output elif [[ "${TEST_CONFIG}" == cachebench ]]; then install_torchaudio install_torchvision @@ -2008,19 +2282,27 @@ elif [[ "${TEST_CONFIG}" == *torchbench* ]]; then LIBTBB_PATH="$(find "$(dirname "$(which python)")/../lib/" -name libtbb.so.12)" export LD_PRELOAD="$LIBTBB_PATH":"$LD_PRELOAD" fi + setup_torch_trace PYTHONPATH=/torchbench test_dynamo_benchmark torchbench "$id" + collect_tlparse_output fi elif [[ "${TEST_CONFIG}" == *inductor_cpp_wrapper* ]]; then install_torchvision + setup_torch_trace PYTHONPATH=/torchbench test_inductor_cpp_wrapper_shard "$SHARD_NUMBER" if [[ "$SHARD_NUMBER" -eq "1" ]]; then test_inductor_aoti_cpp fi + collect_tlparse_output elif [[ "${TEST_CONFIG}" == *inductor_core* ]]; then + setup_torch_trace test_inductor_core + collect_tlparse_output elif [[ "${TEST_CONFIG}" == *inductor* ]]; then install_torchvision + setup_torch_trace test_inductor_shard "${SHARD_NUMBER}" + collect_tlparse_output elif [[ "${TEST_CONFIG}" == *einops* ]]; then test_einops elif [[ "${TEST_CONFIG}" == *dynamo_core* ]]; then diff --git a/.ci/pytorch/test_example_code/CMakeLists.txt b/.ci/pytorch/test_example_code/CMakeLists.txt index 0af9ba6bb655d..c89a728f27649 100644 --- a/.ci/pytorch/test_example_code/CMakeLists.txt +++ b/.ci/pytorch/test_example_code/CMakeLists.txt @@ -8,7 +8,7 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}") add_executable(simple-torch-test simple-torch-test.cpp) target_include_directories(simple-torch-test PRIVATE ${TORCH_INCLUDE_DIRS}) target_link_libraries(simple-torch-test "${TORCH_LIBRARIES}") -set_property(TARGET simple-torch-test PROPERTY CXX_STANDARD 17) +set_property(TARGET simple-torch-test PROPERTY CXX_STANDARD 20) find_package(CUDAToolkit 11.8) diff --git a/.ci/pytorch/windows/arm64/build_libtorch.bat b/.ci/pytorch/windows/arm64/build_libtorch.bat index 1ac14ff697730..33080af122c73 100644 --- a/.ci/pytorch/windows/arm64/build_libtorch.bat +++ b/.ci/pytorch/windows/arm64/build_libtorch.bat @@ -27,8 +27,8 @@ where cl.exe :: change to source directory cd %PYTORCH_ROOT% -:: copy libuv.dll -copy %libuv_ROOT%\lib\Release\uv.dll torch\lib\uv.dll +:: copy libuv.dll (cmake installs the dll to bin/, not lib/Release/) +copy %libuv_ROOT%\bin\uv.dll torch\lib\uv.dll :: create virtual environment python -m venv .venv diff --git a/.ci/pytorch/windows/arm64/build_pytorch.bat b/.ci/pytorch/windows/arm64/build_pytorch.bat index b5c2ef65b84ad..7d10b26339d25 100644 --- a/.ci/pytorch/windows/arm64/build_pytorch.bat +++ b/.ci/pytorch/windows/arm64/build_pytorch.bat @@ -5,6 +5,7 @@ set CMAKE_BUILD_TYPE=%BUILD_TYPE% set CMAKE_C_COMPILER_LAUNCHER=sccache set CMAKE_CXX_COMPILER_LAUNCHER=sccache set libuv_ROOT=%DEPENDENCIES_DIR%\libuv\install +set INSTALL_TEST=0 set MSSdk=1 if defined PYTORCH_BUILD_VERSION ( set PYTORCH_BUILD_VERSION=%PYTORCH_BUILD_VERSION% @@ -27,8 +28,8 @@ where cl.exe :: change to source directory cd %PYTORCH_ROOT% -:: copy libuv.dll -copy %libuv_ROOT%\lib\Release\uv.dll torch\lib\uv.dll +:: copy libuv.dll (cmake installs the dll to bin/, not lib/Release/) +copy %libuv_ROOT%\bin\uv.dll torch\lib\uv.dll :: create virtual environment python -m venv .venv @@ -57,4 +58,4 @@ sccache --show-stats if %errorlevel% neq 0 ( echo "Failed on build_pytorch. (exitcode = %errorlevel%)" exit /b 1 -) \ No newline at end of file +) diff --git a/.ci/pytorch/windows/internal/cuda_config.bat b/.ci/pytorch/windows/internal/cuda_config.bat index 352eb8a3391bf..69059e2e46ee1 100644 --- a/.ci/pytorch/windows/internal/cuda_config.bat +++ b/.ci/pytorch/windows/internal/cuda_config.bat @@ -22,6 +22,10 @@ if "%CUDA_VER%"=="126" ( set "CUDA_DOTTED_VERSION=13.0" set "CUDA_ARCH_LIST=7.5;8.0;8.6;9.0;10.0;12.0" set "VISION_GENCODE=-gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_100,code=compute_100 -gencode=arch=compute_120,code=compute_120" +) else if "%CUDA_VER%"=="132" ( + set "CUDA_DOTTED_VERSION=13.2" + set "CUDA_ARCH_LIST=7.5;8.0;8.6;9.0;10.0;12.0" + set "VISION_GENCODE=-gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_100,code=compute_100 -gencode=arch=compute_120,code=compute_120" ) else ( echo Unknown CUDA version: %CUDA_VER% exit /b 1 diff --git a/.ci/wheel/build_wheel.sh b/.ci/wheel/build_wheel.sh index 6fb63c361f018..2563d5ba31765 100755 --- a/.ci/wheel/build_wheel.sh +++ b/.ci/wheel/build_wheel.sh @@ -97,7 +97,7 @@ fi whl_tmp_dir="${MAC_PACKAGE_WORK_DIR}/dist" mkdir -p "$whl_tmp_dir" -mac_version='macosx-11.0-arm64' +mac_version='macosx-14.0-arm64' libtorch_arch='arm64' # Create a consistent wheel package name to rename the wheel to @@ -125,27 +125,19 @@ popd export TH_BINARY_BUILD=1 export INSTALL_TEST=0 # dont install test binaries into site-packages -export MACOSX_DEPLOYMENT_TARGET=11.0 +export MACOSX_DEPLOYMENT_TARGET=14.0 EXTRA_CONDA_INSTALL_FLAGS="" CONDA_ENV_CREATE_FLAGS="" RENAME_WHEEL=false VERIFY_WHEELNAME=true case $desired_python in - 3.14t) - echo "Using 3.14 deps" - NUMPY_PINNED_VERSION="==2.1.0" - ;; - 3.14) - echo "Using 3.14t deps" - NUMPY_PINNED_VERSION="==2.1.0" - ;; - 3.13t) - echo "Using 3.13t deps" - NUMPY_PINNED_VERSION="==2.1.0" + 3.14*) + echo "Using ${desired_python} deps" + NUMPY_PINNED_VERSION="==2.3.4" ;; - 3.13) - echo "Using 3.13 deps" + 3.13*) + echo "Using ${desired_python} deps" NUMPY_PINNED_VERSION="==2.1.0" ;; 3.12) @@ -179,6 +171,7 @@ retry pip install "${PINNED_PACKAGES[@]}" -r "${pytorch_rootdir}/requirements.tx if [[ -d "/opt/llvm-openmp" ]]; then export OMP_PREFIX=/opt/llvm-openmp else + echo "libomp not found, installing via brew" retry brew install libomp fi diff --git a/.claude/skills/document-public-apis/SKILL.md b/.claude/skills/document-public-apis/SKILL.md index 51cf5c2c34241..211b444fc5461 100644 --- a/.claude/skills/document-public-apis/SKILL.md +++ b/.claude/skills/document-public-apis/SKILL.md @@ -7,7 +7,9 @@ description: Document undocumented public APIs in PyTorch by removing functions This skill documents undocumented public APIs in PyTorch by removing entries from the coverage ignore lists in `docs/source/conf.py` and adding Sphinx autodoc directives (e.g., `autosummary`, `currentmodule`, `autoclass`, `automodule`) to the corresponding `.md` or `.rst` doc source files in `docs/source/`. -**"Documenting" means adding autodoc directives to doc source files — NEVER modifying Python source code.** Do not add or edit docstrings in `.py` files. Do not read or inspect Python source files. Sphinx will pull whatever docstring exists (or render an empty entry if none exists). Your only job is to add the correct directive to the correct doc file. +**"Documenting" means adding autodoc directives to doc source files — NEVER modifying Python source code.** Do not add or edit docstrings in `.py` files. Your only job is to add the correct directive to the correct doc file. + +**IMPORTANT: Before adding any function to the sphinx doctree, verify it has a real docstring.** Use a quick Python check (e.g., `python -c "from torch.module import func; print(bool(func.__doc__))"`) to confirm the function has actual documentation content — not just an empty docstring or a bare `.. warning:: This API is experimental` stub. Functions without meaningful docstrings should be left in the `coverage_ignore_functions`/`coverage_ignore_classes` lists. Adding undocumented functions to the doctree creates empty or near-empty pages that degrade documentation quality. ## Overview @@ -71,9 +73,24 @@ Work through the lists top-to-bottom. Choose enough groups to make meaningful pr If a module group has a **mix** of regular entries and entries with inline comments, still process the group — but only comment out the regular entries. Leave entries with inline comments untouched in the ignore list. +### Step 1b: Verify functions have actual docstrings + +For each function selected in Step 1, check that it has a meaningful docstring by running: + +```bash +python -c "from torch.module.path import func_name; doc = func_name.__doc__; print('HAS DOC' if doc and len(doc.strip()) > 80 else 'NO DOC'); print(repr(doc[:120]) if doc else 'None')" +``` + +A function has a **meaningful docstring** if it has real descriptive content — not just: +- `None` or empty string +- Only a `.. warning:: This API is experimental` stub with no description +- Only a one-line auto-generated signature + +**Functions without meaningful docstrings must stay in the ignore list.** Remove them from your batch. If an entire module group has no functions with docstrings, skip the whole group. + ### Step 2: Present the batch to the user -**Before making any edits**, present the selected module groups and their functions to the user. Show them organized by module: +**Before making any edits**, present the selected module groups and their functions to the user. Indicate which functions passed the docstring check and which were excluded. Show them organized by module: ``` Module: torch.ao.quantization.fx.convert @@ -315,8 +332,8 @@ Also delete any module label comments that no longer have active entries beneath ## Important notes -- **Follow the steps exactly as written.** Do not add extra investigation steps like importing Python modules to check docstrings, inspecting source code to verify function signatures, or running any commands not specified in the instructions. The `make coverage` step is the only verification needed — let it tell you what's wrong. -- **Never modify Python source files (`.py`).** This skill only edits `docs/source/conf.py` and doc source files (`.md`/`.rst`) in `docs/source/`. Do not add or edit docstrings, do not read Python source to check function signatures, do not inspect implementations. +- **Follow the steps exactly as written.** The `make coverage` step is the primary verification for correct Sphinx directives, and Step 1b's docstring check ensures you only document functions that have real content. +- **Never modify Python source files (`.py`).** This skill only edits `docs/source/conf.py` and doc source files (`.md`/`.rst`) in `docs/source/`. Do not add or edit docstrings. The only reason to inspect Python modules is in Step 1b to check whether a docstring exists — never to modify source code. - Entries are commented out in Step 3, verified in Step 6, and cleaned up in Step 8 after verification passes. Never delete uncommented entries directly. - **Read inline comments** on entries before deciding to document them. Entries marked `# deprecated`, `# documented as ...`, `# looks unintentionally public`, or `# legacy helper` should stay in the ignore list. - The `coverage_ignore_functions` list uses bare function names (not fully qualified), so the same name can appear multiple times for different modules. Use the module label comment above each entry to identify which module it belongs to. Be careful during Step 8 cleanup to only delete the correct commented-out lines — commented-out string entries have **quotes** (`# "func_name",`), module label comments do not. diff --git a/.claude/skills/pr-review/SKILL.md b/.claude/skills/pr-review/SKILL.md index ff9cefbc36ced..01d68fb6a7610 100644 --- a/.claude/skills/pr-review/SKILL.md +++ b/.claude/skills/pr-review/SKILL.md @@ -5,7 +5,7 @@ description: Review PyTorch pull requests for code quality, test coverage, secur # PyTorch PR Review Skill -Review PyTorch pull requests focusing on what CI cannot check: code quality, test coverage adequacy, security vulnerabilities, and backward compatibility. Linting, formatting, type checking, and import ordering are handled by CI. +Review PyTorch pull requests focusing on what CI cannot check: code quality, test coverage adequacy, security vulnerabilities, and backward compatibility. ## Usage Modes @@ -115,90 +115,40 @@ only the diff and commit log need to be fetched via git. A single line of code can have deep cross-cutting implications: a missing device guard causes silent data corruption on multi-GPU, a missing `Composite` dispatch key breaks every out-of-tree backend, a manual dtype check instead of `TensorIterator` silently skips type promotion. **Treat every line as potentially load-bearing.** -Do not skim. Do not summarize the diff and move on. Read every changed line and ask: *does this interact with existing PyTorch infrastructure that the author may not know about?* When uncertain, **investigate** — spawn a sub-agent to read the surrounding code, the infrastructure the PR should be using, or the tests that should exist. The cost of a false negative (missing a real issue) is much higher than the cost of investigation. +1. **Investigate, don't guess** — When uncertain whether a checklist item applies, spawn a sub-agent to read the relevant code. A reviewer who guesses wrong provides negative value. +2. **Review the design, not just the implementation** — A PR can have perfectly correct implementation of a bad design. Question side-channel communication, on/off private flags, and demand concrete interface documentation for new contracts between components. +3. **Focus on what CI cannot check** — Don't comment on formatting, linting, type errors, or CI failures. Focus on design quality, interface correctness, thread safety, BC implications, test adequacy, and pattern adherence. +4. **Everything is a must-fix** — There are no "nits." If it's worth mentioning, it's worth fixing. Every inconsistency degrades the codebase over time. +5. **Be specific and actionable** — Reference file paths and line numbers. Name the function/class/file the author should use. +6. **Match the immediate context** — Read how similar features are already implemented in the same file. Pattern mismatches within a file are always wrong. +7. **Assume competence** — The author knows PyTorch; explain only non-obvious context. +8. **No repetition** — Each observation appears in exactly one section of the review output. -## Review Workflow - -### Step 1: Fetch PR Information +### Using sub-agents -**Local CLI mode**: Use `gh` commands to get PR metadata, changed files, full diff, -existing comments/reviews, and associated issue information. +The review checklist is large. You cannot hold the full context of every infrastructure system in your head. **Spawn sub-agents** to investigate whether checklist items apply: read surrounding code, infrastructure the PR should be using, or tests that should exist. Spawn them in parallel for independent areas. A typical medium PR should spawn 3-8 sub-agents. -**Local Branch mode**: Use `git diff` and `git log` against `main` as shown in the -Local Branch Mode section above. - -**GitHub Actions mode**: PR metadata, comments, and reviews are already in the prompt. -Use `git diff origin/...HEAD` for the full diff and -`git log origin/..HEAD --oneline` for the commit log. +## Review Workflow -### Step 2: Understand Context +### Step 1: Understand Context Before reviewing, build understanding of what the PR touches and why: 1. Identify the purpose of the change from title/description/issue 2. Group changes by type (new code, tests, config, docs) 3. Note the scope of changes (files affected, lines changed) -4. **Spawn sub-agents to read the unchanged code surrounding each changed file.** The diff alone is not enough — you need to understand the existing patterns, base classes, and infrastructure in the files being modified. For each significantly changed file, a sub-agent should read the full file (or the relevant class/function) and report back: what patterns does this file follow? What infrastructure does it use? What invariants does it maintain? - -### Step 3: Deep Review — Line-by-Line with Investigation - -This is the core of the review. Go through **every changed line** in the diff and evaluate it against the review checklist in [review-checklist.md](review-checklist.md). - -**How to use sub-agents during review:** +4. Spawn sub-agents to read the unchanged code surrounding each significantly changed file to understand existing patterns and infrastructure -The checklist is large. You cannot hold the full context of every infrastructure system in your head. Instead, when you encounter a changed line that touches a checklist area, **spawn a sub-agent** to investigate whether the checklist item applies. For example: +### Step 2: Deep Review -- A PR adds a new C++ kernel → spawn a sub-agent to check: Does it use TensorIterator? DispatchStub? Structured kernels? AT_DISPATCH? Does it have a meta implementation? A Composite fallback? -- A PR adds a new test → spawn a sub-agent to check: Does an OpInfo exist for this op? Is the test device-generic? Does it use make_tensor, @dtypes, TestCase? -- A PR modifies autograd code → spawn a sub-agent to check: Is derivatives.yaml the right place? Does it use setup_context? Does it have gradcheck tests? -- A PR adds a new operator → spawn a sub-agent to check: Is it in native_functions.yaml? Does it have proper tags? A Composite dispatch? Meta/fake impls? Schema annotations? +Go through **every changed line** in the diff and evaluate it against the review checklist in [review-checklist.md](review-checklist.md). -**Spawn sub-agents in parallel** for independent investigation areas. A typical review of a medium PR should spawn 3-8 sub-agents. Large PRs touching multiple subsystems may need more. +### Step 3: Check Backward Compatibility -**Checklist areas** (see [review-checklist.md](review-checklist.md) for full details): -- Code quality and design -- PyTorch infrastructure — C++ kernels (TensorIterator, DispatchStub, AT_DISPATCH, device guards), CUDA/device management, operator registration and codegen (native_functions.yaml, Composite dispatch, meta/fake implementations), autograd (derivatives.yaml, autograd.Function, gradcheck), Python utilities (pytree, __torch_function__, logging), nn module patterns, Dynamo/Inductor/compile, FX/export, type promotion, serialization, distributed, tensor subclasses -- Testing adequacy (OpInfo, ModuleInfo, device-generic tests, @dtypes, @parametrize, make_tensor) -- Security considerations -- Thread safety and concurrency (Python, C++, CPython C API, NoGIL) -- Performance implications -- Any behavior change not expected by author +Evaluate BC implications per [bc-guidelines.md](bc-guidelines.md). For non-trivial BC questions, spawn a sub-agent to search for existing callers of the modified API. -### Step 4: Check Backward Compatibility +### Step 4: Formulate Review -Evaluate BC implications. See [bc-guidelines.md](bc-guidelines.md) for: -- What constitutes a BC-breaking change -- Required deprecation patterns -- Common BC pitfalls - -For non-trivial BC questions (e.g., "does changing this default break downstream users?"), spawn a sub-agent to search for existing callers of the modified API. - -### Step 5: Formulate Review - -Structure your review with actionable feedback organized by category. Every finding should be traceable to a specific line in the diff and a specific checklist item. - -## Review Areas - -| Area | Focus | Reference | -|------|-------|-----------| -| Code Quality | Abstractions, patterns, complexity | [review-checklist.md](review-checklist.md) | -| API Design | New patterns, flag-based access, broader implications | [review-checklist.md](review-checklist.md) | -| C++ Kernels | TensorIterator, DispatchStub, AT_DISPATCH, structured kernels, device guards, memory format | [review-checklist.md](review-checklist.md) | -| CUDA/Device | C10_CUDA_CHECK, stream/event guards, recordStream, CUDA graphs, AcceleratorHooks | [review-checklist.md](review-checklist.md) | -| Op Registration | native_functions.yaml, Composite fallback, meta/fake impls, tags, schema annotations | [review-checklist.md](review-checklist.md) | -| Autograd | derivatives.yaml, autograd.Function patterns, gradcheck, forward-mode AD, vmap | [review-checklist.md](review-checklist.md) | -| Python Utils | __torch_function__, pytree, logging, deprecation, backends context | [review-checklist.md](review-checklist.md) | -| nn Modules | ModuleList/Dict, nn.init, parametrize, state_dict versioning, LazyModule | [review-checklist.md](review-checklist.md) | -| Dynamo/Inductor | @register_lowering, decompositions, CustomGraphPass, config.patch, graph breaks | [review-checklist.md](review-checklist.md) | -| FX/Export | PassBase, PassManager, Interpreter, subgraph rewriter, ShapeProp, make_fx | [review-checklist.md](review-checklist.md) | -| Type Promotion | elementwise_dtypes, TensorIterator dtype handling, result_type, promoteTypes | [review-checklist.md](review-checklist.md) | -| Serialization | weights_only, safe_globals, skip_data | [review-checklist.md](review-checklist.md) | -| Distributed | DeviceMesh, distributed testing with MultiThreadedPG | [review-checklist.md](review-checklist.md) | -| Tensor Subclasses | _make_wrapper_subclass, __tensor_flatten__/__unflatten__ | [review-checklist.md](review-checklist.md) | -| Testing | OpInfo, ModuleInfo, device-generic, @dtypes, @parametrize, make_tensor | [review-checklist.md](review-checklist.md) | -| Security | Injection, credentials, input handling | [review-checklist.md](review-checklist.md) | -| Performance | Regressions, device handling, memory, profiling, benchmarking | [review-checklist.md](review-checklist.md) | -| Thread Safety | Data races, GIL assumptions, NoGIL, CPython C API | [review-checklist.md](review-checklist.md) | -| BC | Breaking changes, deprecation | [bc-guidelines.md](bc-guidelines.md) | +Structure your review with actionable feedback organized by category. Every finding should be traceable to a specific line in the diff. ## Output Format @@ -258,20 +208,9 @@ When requested, add file-specific feedback with line references: - `torch/nn/modules/linear.py:78` - This allocation could be moved outside the loop ``` -## Key Principles - -1. **Investigate, don't guess** - When uncertain whether a checklist item applies, spawn a sub-agent to read the relevant infrastructure code. A reviewer who guesses wrong provides negative value. A reviewer who investigates and reports findings provides immense value. -2. **Every line matters** - A single missing `C10_CUDA_KERNEL_LAUNCH_CHECK()`, a single `weights_only=False`, a single missing Composite dispatch key — each of these is a real bug that affects real users. Do not skip lines. -3. **No repetition** - Each observation appears in exactly one section. Never repeat the same issue, concern, or suggestion across multiple sections. If an issue spans categories (e.g., a security issue that also affects performance), place it in the most relevant section only. -4. **Focus on what CI cannot check** - Don't comment on formatting, linting, or type errors -5. **Be specific** - Reference file paths and line numbers. Every finding should point to a concrete line in the diff. -6. **Be actionable** - Provide concrete suggestions with the right infrastructure to use, not vague concerns. If flagging a missing pattern, name the function/class/file the author should use. -7. **Be proportionate** - Minor issues shouldn't block, but note them -8. **Assume competence** - The author knows PyTorch; explain only non-obvious context. The value of this review is in catching infrastructure patterns the author may not know about, not in explaining basic programming. - ## Files to Reference -When reviewing, consult these project files for context. **Spawn sub-agents to read these** rather than relying on memory — the files change frequently: +When reviewing, consult these project files for context — read them rather than relying on memory, as they change frequently: - `CLAUDE.md` - Coding style philosophy and testing patterns - `CONTRIBUTING.md` - PR requirements and review process - `torch/testing/_internal/common_utils.py` - Test patterns and utilities diff --git a/.claude/skills/pr-review/review-checklist.md b/.claude/skills/pr-review/review-checklist.md index 2151f2977b2be..7c588027cbdca 100644 --- a/.claude/skills/pr-review/review-checklist.md +++ b/.claude/skills/pr-review/review-checklist.md @@ -7,7 +7,10 @@ This checklist covers areas that CI cannot check. Skip items related to linting, ### Abstractions and Design - [ ] **Clear abstractions** - State management is explicit; no dynamic attribute setting/getting -- [ ] **Match existing patterns** - Code follows architectural patterns already in the codebase +- [ ] **No side-channel communication** - If behavior changes based on a hidden flag or dynamically-set attribute, the interface itself should change instead (different function signature, different class, different code path). Side-channel patterns (set a private flag in one place, check it in another via `getattr`) create undocumented behavioral modes +- [ ] **Proper interface, not on/off flags** - A private boolean that switches between two fundamentally different behaviors should be two separate code paths or a proper interface change, not a flag +- [ ] **Interface documentation** - New internal calling conventions, protocols, or contracts between components must have concrete documentation: what the caller provides, what the callee receives, what invariants hold, and cleanup responsibilities. Motivational comments ("this allows X") are not interface documentation +- [ ] **Match existing patterns in the same file** - Before accepting new code in a file, read how similar features are already implemented in that same file. If the file uses class attributes for boolean flags, new boolean flags must use class attributes. If the file uses a specific setter pattern, new setters must use the same pattern - [ ] **No over-engineering** - Only requested changes are made; no speculative features - [ ] **No premature abstraction** - Helpers and utilities are only created when reused; three similar lines is better than a one-use helper - [ ] **No trivial helpers** - Avoid 1-2 LOC helper functions used only once (unless significantly improves readability) @@ -36,17 +39,9 @@ When a PR introduces new API patterns, carefully evaluate the broader implicatio - [ ] **No fragile init ordering** - If multiple imports/calls must happen in a specific undocumented order, flag the design. Dependencies should be explicit or combined into a single entry point - [ ] **Idempotent global state** - Registries and global lists that accumulate entries must handle multiple calls safely (no duplicate registration, clear cleanup story) -### Common Issues to Flag - -- Dynamic `setattr`/`getattr` for state management (prefer explicit class members) -- Unused imports, variables, or dead code paths -- Copy-pasted code that could be a shared helper -- Magic numbers without explanation -- Overly defensive error handling for impossible cases - ## PyTorch Infrastructure -When a PR touches code in the scope of any item below, **stop and investigate** whether the established infrastructure should be used. Spawn a sub-agent to read the relevant infrastructure code and determine if the PR should be using it instead of rolling its own solution. +When a PR touches code in the scope of any item below, **stop and investigate** whether the established infrastructure should be used. ### C++ Kernel Infrastructure @@ -55,8 +50,9 @@ When a PR touches code in the scope of any item below, **stop and investigate** - [ ] **Structured Kernels** — PR adds a new ATen operator with separate hand-written functional, inplace, and out= variants instead of using `structured: True` + `structured_delegate` in `native_functions.yaml` to generate boilerplate - [ ] **TORCH_CHECK variants** — PR uses generic `TORCH_CHECK` for conditions that have a more specific variant: `ValueError` → `TORCH_CHECK_VALUE`, `IndexError` → `TORCH_CHECK_INDEX`, `TypeError` → `TORCH_CHECK_TYPE`, `NotImplementedError` → `TORCH_CHECK_NOT_IMPLEMENTED` - [ ] **AT_DISPATCH macros** — PR manually switches on `dtype` with `if (dtype == kFloat) ... else if (dtype == kDouble)` instead of using `AT_DISPATCH_FLOATING_TYPES`, `AT_DISPATCH_ALL_TYPES_AND`, or the `AT_DISPATCH_SWITCH` / `AT_DISPATCH_CASE` pattern from `aten/src/ATen/Dispatch.h` -- [ ] **Device guards (RAII)** — PR manually saves/restores device context (`cudaSetDevice` + try/catch) instead of using `DeviceGuard` or `OptionalDeviceGuard` from `c10/core/DeviceGuard.h` +- [ ] **Device guards (RAII)** — PR manually saves/restores device context (`cudaSetDevice` + try/catch) instead of using `DeviceGuard` or `OptionalDeviceGuard` from `c10/core/DeviceGuard.h`. **Note:** Operators registered in `native_functions.yaml` get automatic `DeviceGuard` insertion from codegen (controlled by `device_guard: True`, the default) — do NOT flag missing device guards for these ops unless they explicitly set `device_guard: False` - [ ] **Memory format propagation** — PR allocates output tensors with `at::empty(shape, options)` (defaulting to contiguous) without calling `input.suggest_memory_format()` to preserve ChannelsLast or other input formats +- [ ] **Subclass-safe tensor allocation** — PR uses `at::empty(shape, input.options())` instead of `input.new_empty(shape)` or `at::empty_like(input)`, which don't propagate tensor subclass metadata - [ ] **TORCH_LIBRARY operator registration** — PR registers operators using manual dispatcher calls instead of `TORCH_LIBRARY` / `TORCH_LIBRARY_IMPL` macros from `torch/library.h` - [ ] **TORCH_WARN_DEPRECATION** — PR uses `TORCH_WARN` for deprecation notices instead of `TORCH_WARN_DEPRECATION` which issues a proper `DeprecationWarning` @@ -189,32 +185,9 @@ When a PR touches code in the scope of any item below, **stop and investigate** ### Test Quality - [ ] **Edge cases covered** - Tests include boundary conditions, empty inputs, error cases -- [ ] **Error conditions tested** - Expected exceptions are tested with `assertRaises` or `assertRaisesRegex` -- [ ] **No duplicated test logic** - Similar tests share a private helper method (e.g., `_test_foo(config)`) called from individual tests with different configs - -**Example of good test structure:** -```python -def _test_feature_with_config(self, flag, expected_shape): - """Shared test logic called by device-specific tests.""" - x = torch.randn(10) - result = my_feature(x, flag) - self.assertEqual(result.shape, expected_shape) - -def test_feature_enabled(self): - self._test_feature_with_config(True, (10, 10)) - -def test_feature_disabled(self): - self._test_feature_with_config(False, (10, 5)) -``` - -### Common Testing Issues - -- Tests that only check the happy path without error cases -- Duplicated test code that should be a parameterized helper -- Manual operator tests that duplicate existing OpInfo coverage — the fix should update the OpInfo's dtype list instead -- Tests that don't clean up resources (files, CUDA memory) -- Flaky tests (timing-dependent, order-dependent, golden value) -- Tests that skip without clear justification +- [ ] **Error conditions tested** - Expected exceptions are tested with `assertRaisesRegex`, not bare `assertRaises`. `assertRaisesRegex` verifies both the exception type and message, catching cases where the right exception is raised for the wrong reason. Bare `assertRaises` should be flagged — always require a message pattern match +- [ ] **No duplicated test logic** - Similar tests share a private helper method called from individual tests with different configs +- [ ] **Use weakref for lifetime testing** - PR uses `sys.getrefcount()` to test whether objects are kept alive. Use `weakref.ref()` instead — create a weak reference, delete the strong references, then check if the weakref is dead (`wr() is None`). `sys.getrefcount` is a CPython implementation detail that varies across versions and is fragile ## Security @@ -258,7 +231,7 @@ This is particularly important for PyTorch's autograd, which has multi-threaded - [ ] **GIL held for Python object access** - Any code that touches `PyObject*` (incref, decref, attribute access, container mutation) must hold the GIL. When releasing the GIL for long-running C++ work (`Py_BEGIN_ALLOW_THREADS`), verify no Python objects are accessed in that region - [ ] **Borrowed references across GIL release** - Borrowed references (`PyTuple_GET_ITEM`, `PyList_GET_ITEM`) become unsafe if the GIL is released and reacquired, since another thread may have mutated the container -- [ ] **Decref-before-update hazard** - When replacing an item in a container (tuple, list, dict), update the container slot first, then `Py_DECREF` the old value. Decref can trigger `__del__` finalizers that re-enter and observe the container in an inconsistent state. Without the GIL (free-threaded builds), this is also a data race. +- [ ] **Decref-before-update hazard** - When replacing an item in a container (tuple, list, dict), update the container slot first, then `Py_DECREF` the old value. Decref can trigger `__del__` finalizers that re-enter and observe the container in an inconsistent state. Without the GIL (free-threaded builds), this is also a data race. This is **always** a must-fix — even if "safe in practice" because of refcount guarantees, the pattern is wrong and breaks under NoGIL. The correct pattern costs nothing extra ### Free-Threaded Python (NoGIL, PEP 703) @@ -298,10 +271,3 @@ CPython 3.13t+ can run without the GIL. Code that was previously safe under the - [ ] **Use torch.profiler** - PR adds manual `time.time()` instrumentation instead of using `torch.profiler.profile()` context manager with `schedule()` and `tensorboard_trace_handler()` - [ ] **Use torch.utils.benchmark.Timer** - PR benchmarks with `time.time()` loops instead of `torch.utils.benchmark.Timer` which handles warmup, statistics, and proper CUDA synchronization - -### Common Performance Issues - -- Creating new tensors inside training loops instead of pre-allocating -- Synchronous CUDA operations where async would work -- Keeping computation graph alive longer than needed -- Redundant clones or copies diff --git a/.claude/skills/pt2-bug-basher/SKILL.md b/.claude/skills/pt2-bug-basher/SKILL.md index 8de7587d56f99..5c30f9002f79c 100644 --- a/.claude/skills/pt2-bug-basher/SKILL.md +++ b/.claude/skills/pt2-bug-basher/SKILL.md @@ -10,17 +10,19 @@ Debug test failures and runtime errors in the PyTorch 2 compiler stack (Dynamo, ## Workflow Summary -1. **Reproduce** -- Get a consistent reproduction of the failure -2. **Minimize** -- Reduce the repro to the smallest possible standalone case. Strip away unrelated model logic, use minimal tensor shapes, and isolate the specific op or pattern that triggers the bug. -3. **Add a unit test** -- **Do this BEFORE diving into code search or root cause investigation.** Add a failing test to the codebase that captures the bug. Place it in a specific, topic-appropriate test file (e.g., `test/dynamo/test_repros.py`, `test/inductor/test_torchinductor.py`, `test/export/test_export.py`). **Avoid `test/dynamo/test_misc.py`** — it is already oversized; find a more specific test file that matches the area of the bug. Use `torch.testing._internal.common_utils.TestCase` and `run_tests`. The test must fail before the fix and pass after. Having the test first keeps you grounded — you know exactly what "fixed" looks like before you start exploring the codebase. -4. **Gather logs** -- Run with appropriate `TORCH_LOGS` settings -5. **Classify** -- Use the [Error Triage](#error-triage) table to identify the category -6. **Inspect artifacts** -- Check FX graphs, IR, and generated code via `TORCH_COMPILE_DEBUG=1` -7. **Identify root cause** -- Trace from the error back through the compilation pipeline -8. **Fix** -- Apply the fix -9. **Verify** -- Run the new unit test AND nearby related existing tests (e.g., if you changed how `is_exporting` works, also run the existing `test_is_exporting` export test). Use `pytest -k` to quickly run related tests by name. The task is not complete until all pass. -10. **Self-review** -- Use the `/pr-review` skill to review your own changes before presenting them. Fix any issues it flags. -11. **Celebrate** -- Summarize the changes: explain the root cause, what was changed and why, and which tests were added/verified. Then tell the user the bug is squashed. Include a fun, varied motivational message or easter egg to keep spirits high (e.g., a pun, a quote, an ASCII art bug getting squashed). Keep it short and different each time. +1. **Environment check** -- Ask the user which conda environment to use. Verify it is active by checking `$CONDA_DEFAULT_ENV`. Then run `python -c "import torch; print(torch.__version__)"` to confirm torch is importable and report the version. If the environment is not active or torch cannot be imported, stop and ask the user to activate the correct environment before proceeding. +2. **Reproduce** -- Get a consistent reproduction of the failure +3. **Minimize** -- Reduce the repro to the smallest possible standalone case. Strip away unrelated model logic, use minimal tensor shapes, and isolate the specific op or pattern that triggers the bug. +4. **Add a unit test** -- **Do this BEFORE diving into code search or root cause investigation.** Add a failing test to the codebase that captures the bug. Place it in a specific, topic-appropriate test file (e.g., `test/dynamo/test_repros.py`, `test/inductor/test_torchinductor.py`, `test/export/test_export.py`). **Avoid `test/dynamo/test_misc.py`** — it is already oversized; find a more specific test file that matches the area of the bug. Use `torch.testing._internal.common_utils.TestCase` and `run_tests`. The test must fail before the fix and pass after. Having the test first keeps you grounded — you know exactly what "fixed" looks like before you start exploring the codebase. +5. **Validate on main** -- Use `EnterWorktree` to create a worktree checked out at `main`. Copy the new test file into the worktree and run the test there to confirm it **fails** on main. If the test passes on main, stop — the test may not be capturing the right bug, or the bug may already be fixed. Exit the worktree with `ExitWorktree` (action: remove) and return to the working branch before continuing. +6. **Gather logs** -- Run with appropriate `TORCH_LOGS` settings +7. **Classify** -- Use the [Error Triage](#error-triage) table to identify the category +8. **Inspect artifacts** -- Check FX graphs, IR, and generated code via `TORCH_COMPILE_DEBUG=1` +9. **Identify root cause** -- Trace from the error back through the compilation pipeline +10. **Fix** -- Apply the fix +11. **Verify** -- Run the new unit test AND nearby related existing tests (e.g., if you changed how `is_exporting` works, also run the existing `test_is_exporting` export test). Use `pytest -k` to quickly run related tests by name. The task is not complete until all pass. +12. **Self-review** -- Use the `/pr-review` skill to review your own changes before presenting them. Fix any issues it flags. +13. **Celebrate** -- Summarize the changes: explain the root cause, what was changed and why, and which tests were added/verified. Then tell the user the bug is squashed. Include a fun, varied motivational message or easter egg to keep spirits high (e.g., a pun, a quote, an ASCII art bug getting squashed). Keep it short and different each time. ## Investigation Strategy diff --git a/.claude/skills/triaging-issues/SKILL.md b/.claude/skills/triaging-issues/SKILL.md index 9cc741540a9f4..e774d6dd40cb6 100644 --- a/.claude/skills/triaging-issues/SKILL.md +++ b/.claude/skills/triaging-issues/SKILL.md @@ -156,7 +156,7 @@ If the issue belongs in another repo (vision/text/audio/RL/ExecuTorch/etc.), tra **PT2 is NOT a redirect.** `oncall: pt2` is not like the other oncall labels in Step 3. PT2 issues continue through Steps 4–7 for full triage — add `oncall: pt2`, then proceed to label with `module:` labels, mark `triaged`, etc. -See [pt2-triage-rubric.md](pt2-triage-rubric.md) for detailed labeling decisions on which `module:` labels to apply. +**Every `oncall: pt2` issue MUST have at least one `module:` label.** The PT2 oncall queue is too broad without a module label — the team needs to know which component is affected (e.g., `module: dynamo`, `module: inductor`, `module: helion`, `module: dynamic shapes`). If you cannot determine the specific module, use `module: compile ux` as a fallback, but always try to be specific first. See [pt2-triage-rubric.md](pt2-triage-rubric.md) for detailed guidance. ### 3) Redirect to Secondary Oncall diff --git a/.claude/skills/triaging-issues/labels.json b/.claude/skills/triaging-issues/labels.json index 7afa211cd5dc1..fa3703186472b 100644 --- a/.claude/skills/triaging-issues/labels.json +++ b/.claude/skills/triaging-issues/labels.json @@ -437,7 +437,7 @@ }, { "name": "module: helion", - "description": "" + "description": "Helion (pytorch/helion) is a high-level Python DSL for writing GPU kernels that compiles to Triton. Users write kernels with @helion.kernel decorator and helion.language (hl.tile(), etc.) using familiar PyTorch operations. Helion kernels integrate with torch.compile/inductor via template fusion hooks. Apply when the issue involves helion kernel authoring, helion compilation, helion-to-Triton codegen, or helion+inductor template fusion. Do NOT apply for plain Triton kernel issues (use module: inductor) or general torch.compile issues unless helion is specifically involved." }, { "name": "module: higher order operators", diff --git a/.claude/skills/triaging-issues/pt2-triage-rubric.md b/.claude/skills/triaging-issues/pt2-triage-rubric.md index 1d9742670eca3..967725380f7be 100644 --- a/.claude/skills/triaging-issues/pt2-triage-rubric.md +++ b/.claude/skills/triaging-issues/pt2-triage-rubric.md @@ -2,6 +2,8 @@ This rubric guides labeling decisions for PT2 oncall triage. +**Every `oncall: pt2` issue MUST have at least one `module:` label.** The PT2 queue is too broad without one — the team needs to know which component is affected. Use the sections below to determine the right module label(s). + ## 1. Component Isolation - Be Precise, Don't Over-Tag ### Dynamo vs Dynamic Shapes @@ -135,7 +137,34 @@ Check for existing labels before inventing categories: --- -## 7. functorch + compile +## 7. Helion Kernel Issues + +[Helion](https://github.com/pytorch/helion) is a high-level Python DSL for writing GPU kernels. Users write kernels with `@helion.kernel` and `helion.language` (`hl.tile()`, etc.) using standard PyTorch ops, and Helion compiles them to Triton. Helion kernels can also be fused into `torch.compile` graphs via inductor's template fusion hooks. + +### Identifying Helion Issues + +| Signal | Label | +|--------|-------| +| Issue mentions `helion`, `@helion.kernel`, `helion.language`, `hl.tile` | `module: helion` | +| Error traceback includes `helion/` or `helion.` frames | `module: helion` | +| Helion kernel produces wrong results (standalone, no torch.compile) | `module: helion` only | +| Helion kernel fails or miscompiles under `torch.compile` | `module: helion` + `module: inductor` | +| Inductor template fusion bug triggered by a Helion template | `module: helion` + `module: inductor` | + +### Routing + +- Helion issues get `oncall: pt2` — Helion is a PT2 component. +- If the bug is purely in Helion's own compilation (standalone kernel, not under torch.compile), apply `module: helion` without `module: inductor`. +- If the bug is in how inductor fuses or emits Helion templates, apply both `module: helion` and `module: inductor`. + +### Common Mistakes + +- **Don't confuse Helion with raw Triton**: If the user is writing Triton kernels directly (using `@triton.jit`, `tl.load`, etc.) without Helion, that's `module: inductor`, not `module: helion`. +- **Don't apply `module: helion` for general inductor codegen bugs**: Just because inductor generates Triton code doesn't make it a Helion issue. Helion is a specific DSL — look for explicit `helion` imports or mentions. + +--- + +## 8. Functorch + Compile | Situation | Labels | |-----------|--------| @@ -144,7 +173,7 @@ Check for existing labels before inventing categories: --- -## 8. High Priority Criteria +## 9. High Priority Criteria **This is critical** You should not explicitly add `high priority` - add `triage review` instead so that it is reviewed at the next triage meeting by the oncall. @@ -162,7 +191,7 @@ Mark `triage review` if ANY of these apply: --- -## 9. Fuzzer Issues +## 10. Fuzzer Issues For `topic: fuzzer` issues: @@ -174,7 +203,7 @@ For `topic: fuzzer` issues: --- -## 10. Quick Label Reference +## 11. Quick Label Reference ### Core Components - `module: dynamo` - Tracing, bytecode, graph breaks @@ -183,6 +212,7 @@ For `topic: fuzzer` issues: - `module: pt2-dispatcher` - AOT autograd, functionalization, FakeTensor - `module: cuda graphs` - CUDA graph capture/replay - `module: flex attention` - Flex attention API +- `module: helion` - Helion DSL kernel authoring, compilation, and inductor fusion ### Holistic Areas - `module: compile ux` - Error messages, APIs, programming model diff --git a/.claude/skills/triaging-issues/scripts/validate_labels.py b/.claude/skills/triaging-issues/scripts/validate_labels.py index cd2b4f22e5098..df489189a20b5 100755 --- a/.claude/skills/triaging-issues/scripts/validate_labels.py +++ b/.claude/skills/triaging-issues/scripts/validate_labels.py @@ -208,12 +208,18 @@ def main(): except json.JSONDecodeError as e: debug_log(f"JSON decode error: {e}") print(f"Hook error: Invalid JSON input: {e}", file=sys.stderr) - print("Hook was unable to validate labels; stopping triage.", file=sys.stderr) + print( + "Hook was unable to validate labels; stopping triage.", + file=sys.stderr, + ) sys.exit(2) except Exception as e: debug_log(f"Unexpected error: {type(e).__name__}: {e}") print(f"Hook error: {e}", file=sys.stderr) - print("Hook was unable to validate labels; stopping triage.", file=sys.stderr) + print( + "Hook was unable to validate labels; stopping triage.", + file=sys.stderr, + ) sys.exit(2) diff --git a/.flake8 b/.flake8 index 12d7432bbb7c6..f5507047d7f39 100644 --- a/.flake8 +++ b/.flake8 @@ -12,13 +12,21 @@ ignore = # to line this up with executable bit EXE001, # these ignores are from flake8-bugbear; please fix! - B007,B008,B017,B019,B023,B028,B903,B905,B906,B907,B908,B910 + B007,B008,B017,B019,B023,B028,B903,B905,B906,B907,B908,B910, # these ignores are from flake8-simplify. please fix or ignore with commented reason SIM105,SIM108,SIM110,SIM111,SIM113,SIM114,SIM115,SIM116,SIM117,SIM118,SIM119,SIM12, # SIM104 is already covered by pyupgrade ruff SIM104, # flake8-simplify code styles - SIM102,SIM103,SIM106,SIM112 + SIM102,SIM103,SIM106,SIM112, + # Codes where ruff is the source of truth. Ruff handles these (sometimes + # differently than flake8). As ruff promotes preview codes, move them here + # from ruff's external list. + B001,B036,B902,B904,B950, + C406,C417,C419,C901, + E121,E122,E124,E128,E131,E712,E722,E723,E731, + F401,F722,F723,F811,F812, + G001,G003,G004,G010,G200,G201,G202 per-file-ignores = __init__.py: F401 test/**: F821 diff --git a/.github/actionlint.yaml b/.github/actionlint.yaml index 91c35a03948b5..b894d7214fced 100644 --- a/.github/actionlint.yaml +++ b/.github/actionlint.yaml @@ -68,10 +68,13 @@ self-hosted-runner: # Org wise AWS `mac2.metal` runners (2020 Mac mini hardware powered by Apple silicon M1 processors) - macos-m1-stable - macos-m1-14 + - macos-m2-15 + - macos-m2-26 # GitHub-hosted MacOS runners - macos-latest-xlarge - macos-13-xlarge - macos-14-xlarge + - macos-26-xlarge # Organization-wide Intel hosted XPU runners - linux.idc.xpu # Organization-wide Google Cloud TPU runners diff --git a/.github/actions/download-build-artifacts/action.yml b/.github/actions/download-build-artifacts/action.yml index c44b6a4083448..ea78440e87686 100644 --- a/.github/actions/download-build-artifacts/action.yml +++ b/.github/actions/download-build-artifacts/action.yml @@ -18,14 +18,14 @@ runs: using: composite steps: - name: Download PyTorch Build Artifacts from S3 - if: ${{ !inputs.use-gha }} + if: inputs.use-gha == '' || inputs.use-gha == 'false' uses: seemethere/download-artifact-s3@v4 with: name: ${{ inputs.name }} s3-bucket: ${{ inputs.s3-bucket }} - name: Download PyTorch Build Artifacts from GHA - if: ${{ inputs.use-gha }} + if: inputs.use-gha != '' && inputs.use-gha != 'false' uses: actions/download-artifact@v4 with: name: ${{ inputs.name }} diff --git a/.github/actions/download-td-artifacts/action.yml b/.github/actions/download-td-artifacts/action.yml index 18766bf670f63..3b25f14ef6520 100644 --- a/.github/actions/download-td-artifacts/action.yml +++ b/.github/actions/download-td-artifacts/action.yml @@ -11,13 +11,13 @@ runs: using: composite steps: - name: Download TD Artifacts from S3 - if: ${{ !inputs.use-gha }} + if: inputs.use-gha == '' || inputs.use-gha == 'false' uses: seemethere/download-artifact-s3@v4 with: name: td_results - name: Download TD Artifacts from GHA - if: inputs.use-gha + if: inputs.use-gha != '' && inputs.use-gha != 'false' uses: actions/download-artifact@v4 with: name: td_results.json diff --git a/.github/actions/linux-test/action.yml b/.github/actions/linux-test/action.yml deleted file mode 100644 index 8ad74e308168a..0000000000000 --- a/.github/actions/linux-test/action.yml +++ /dev/null @@ -1,406 +0,0 @@ -name: linux-test - -inputs: - build-environment: - required: true - type: string - description: Top-level label for what's being built/tested. - test-matrix: - required: true - type: string - description: JSON description of what test configs to run. - docker-image: - required: true - type: string - description: Docker image to run in. - sync-tag: - required: false - type: string - default: "" - description: | - If this is set, our linter will use this to make sure that every other - job with the same `sync-tag` is identical. - use-gha: - required: false - type: string - default: "" - description: If set to any value, upload to GHA. Otherwise upload to S3. - dashboard-tag: - required: false - type: string - default: "" - s3-bucket: - description: S3 bucket to download artifact - required: false - type: string - default: "gha-artifacts" - aws-role-to-assume: - description: role to assume for downloading artifacts - required: false - type: string - default: "" - HUGGING_FACE_HUB_TOKEN: - description: | - HF Auth token to avoid rate limits when downloading models or datasets from hub - required: false - default: "" - GITHUB_TOKEN: - description: GitHub token - required: true - disable-monitor: - description: | - [Experimental] Disable utilization monitoring for tests. - Currently, by default we disable the monitor job and only look for specific tests, - since we are investigating the behaviour of the monitor script with different tests. - required: false - type: boolean - default: true -#env: -# GIT_DEFAULT_BRANCH: ${{ inputs.default_branch }} - -runs: - using: composite - steps: - - name: Setup Linux - uses: ./.github/actions/setup-linux - - - name: Login to ECR - uses: ./.github/actions/ecr-login - with: - aws-role-to-assume: ${{ inputs.aws-role-to-assume }} - - - name: Calculate docker image - id: calculate-docker-image - uses: pytorch/test-infra/.github/actions/calculate-docker-image@main - with: - docker-image-name: ${{ inputs.docker-image }} - - - name: Use following to pull public copy of the image - id: print-ghcr-mirror - env: - ECR_DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }} - shell: bash - run: | - tag=${ECR_DOCKER_IMAGE##*/} - echo "docker pull ghcr.io/pytorch/ci-image:${tag/:/-}" - - - name: Pull docker image - uses: pytorch/test-infra/.github/actions/pull-docker-image@main - with: - docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }} - - - name: Check if in a container runner - shell: bash - id: check_container_runner - run: echo "IN_CONTAINER_RUNNER=$(if [ -f /.inarc ] || [ -f /.incontainer ]; then echo true ; else echo false; fi)" >> "$GITHUB_OUTPUT" - - - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG - id: install-nvidia-driver - uses: pytorch/test-infra/.github/actions/setup-nvidia@main - if: ${{ contains(inputs.build-environment, 'cuda') && !contains(matrix.config, 'nogpu') && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }} - - - name: Setup GPU_FLAG for docker run - id: setup-gpu-flag - run: echo "GPU_FLAG=--gpus all -e NVIDIA_DRIVER_CAPABILITIES=all" >> "${GITHUB_ENV}" - if: ${{ contains(inputs.build-environment, 'cuda') && !contains(matrix.config, 'nogpu') && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'true' }} - - - name: Setup SCCACHE_SERVER_PORT environment for docker run when on container - id: setup-sscache-port-flag - run: echo "SCCACHE_SERVER_PORT_DOCKER_FLAG=-e SCCACHE_SERVER_PORT=$((RUNNER_UID + 4226))" >> "${GITHUB_ENV}" - if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'true' }} - - - name: Lock NVIDIA A100 40GB Frequency - shell: bash - run: | - sudo nvidia-smi -pm 1 - sudo nvidia-smi -ac 1215,1410 - nvidia-smi - if: ${{ contains(matrix.runner, 'a100') && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }} - - - name: Start monitoring script - id: monitor-script - if: ${{ !inputs.disable-monitor }} - shell: bash - continue-on-error: true - run: | - python3 -m pip install psutil==5.9.8 nvidia-ml-py==11.525.84 - python3 -m tools.stats.monitor > usage_log.txt 2>&1 & - echo "monitor-script-pid=${!}" >> "${GITHUB_OUTPUT}" - - - name: Download build artifacts - uses: ./.github/actions/download-build-artifacts - with: - name: ${{ inputs.build-environment }} - s3-bucket: ${{ inputs.s3-bucket }} - - - name: Download TD artifacts - continue-on-error: true - uses: ./.github/actions/download-td-artifacts - - - name: Parse ref - id: parse-ref - shell: bash - run: .github/scripts/parse_ref.py - - - name: Get workflow job id - id: get-job-id - uses: ./.github/actions/get-workflow-job-id - if: always() - with: - github-token: ${{ inputs.GITHUB_TOKEN }} - - - name: Check for keep-going label and re-enabled test issues - # This uses the filter-test-configs action because it conveniently - # checks for labels and re-enabled test issues. It does not actually do - # any filtering. All filtering is done in the build step. - id: keep-going - uses: ./.github/actions/filter-test-configs - with: - github-token: ${{ inputs.GITHUB_TOKEN }} - test-matrix: ${{ inputs.test-matrix }} - job-name: ${{ steps.get-job-id.outputs.job-name }} - - - name: Test - id: test - env: - BUILD_ENVIRONMENT: ${{ inputs.build-environment }} - PR_NUMBER: ${{ github.event.pull_request.number }} - GITHUB_REPOSITORY: ${{ github.repository }} - GITHUB_WORKFLOW: ${{ github.workflow }} - GITHUB_JOB: ${{ github.job }} - GITHUB_RUN_ID: ${{ github.run_id }} - GITHUB_RUN_NUMBER: ${{ github.run_number }} - GITHUB_RUN_ATTEMPT: ${{ github.run_attempt }} - JOB_ID: ${{ steps.get-job-id.outputs.job-id }} - JOB_NAME: ${{ steps.get-job-id.outputs.job-name }} - BRANCH: ${{ steps.parse-ref.outputs.branch }} - SHA1: ${{ github.event.pull_request.head.sha || github.sha }} - BASE_SHA: ${{ github.event.pull_request.base.sha || github.sha }} - TEST_CONFIG: ${{ matrix.config }} - SHARD_NUMBER: ${{ matrix.shard }} - NUM_TEST_SHARDS: ${{ matrix.num_shards }} - REENABLED_ISSUES: ${{ steps.keep-going.outputs.reenabled-issues }} - CONTINUE_THROUGH_ERROR: ${{ steps.keep-going.outputs.keep-going }} - VERBOSE_TEST_LOGS: ${{ steps.keep-going.outputs.ci-verbose-test-logs }} - TEST_SHOWLOCALS: ${{ steps.keep-going.outputs.ci-test-showlocals }} - NO_TEST_TIMEOUT: ${{ steps.keep-going.outputs.ci-no-test-timeout }} - NO_TD: ${{ steps.keep-going.outputs.ci-no-td }} - TD_DISTRIBUTED: ${{ steps.keep-going.outputs.ci-td-distributed }} - SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2 - SCCACHE_REGION: us-east-1 - SCCACHE_S3_KEY_PREFIX: ${{ github.workflow }} - SHM_SIZE: ${{ contains(inputs.build-environment, 'cuda') && '2g' || '1g' }} - DOCKER_IMAGE: ${{ inputs.docker-image }} - XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }} - XLA_CLANG_CACHE_S3_BUCKET_NAME: ossci-compiler-clang-cache-circleci-xla - PYTORCH_TEST_CUDA_MEM_LEAK_CHECK: ${{ matrix.mem_leak_check && '1' || '0' }} - PYTORCH_TEST_RERUN_DISABLED_TESTS: ${{ matrix.rerun_disabled_tests && '1' || '0' }} - DASHBOARD_TAG: ${{ inputs.dashboard-tag }} - HUGGING_FACE_HUB_TOKEN: ${{ inputs.HUGGING_FACE_HUB_TOKEN }} - SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }} - IS_A100_RUNNER: ${{ contains(matrix.runner, 'a100') && '1' || '0' }} - - shell: bash - run: | - set -x - - if [[ $TEST_CONFIG == 'multigpu' ]]; then - TEST_COMMAND=.ci/pytorch/multigpu-test.sh - elif [[ $BUILD_ENVIRONMENT == *onnx* ]]; then - TEST_COMMAND=.ci/onnx/test.sh - else - TEST_COMMAND=.ci/pytorch/test.sh - fi - - # detached container should get cleaned up by teardown_ec2_linux - # TODO: Stop building test binaries as part of the build phase - # Used for GPU_FLAG since that doesn't play nice - # shellcheck disable=SC2086,SC2090 - container_name=$(docker run \ - ${GPU_FLAG:-} \ - ${SCCACHE_SERVER_PORT_DOCKER_FLAG:-} \ - -e BUILD_ENVIRONMENT \ - -e PR_NUMBER \ - -e GITHUB_ACTIONS \ - -e GITHUB_REPOSITORY \ - -e GITHUB_WORKFLOW \ - -e GITHUB_JOB \ - -e GITHUB_RUN_ID \ - -e GITHUB_RUN_NUMBER \ - -e GITHUB_RUN_ATTEMPT \ - -e JOB_ID \ - -e JOB_NAME \ - -e BASE_SHA \ - -e BRANCH \ - -e SHA1 \ - -e AWS_DEFAULT_REGION \ - -e IN_WHEEL_TEST \ - -e SHARD_NUMBER \ - -e TEST_CONFIG \ - -e NUM_TEST_SHARDS \ - -e REENABLED_ISSUES \ - -e CONTINUE_THROUGH_ERROR \ - -e VERBOSE_TEST_LOGS \ - -e NO_TEST_TIMEOUT \ - -e NO_TD \ - -e TD_DISTRIBUTED \ - -e PR_LABELS \ - -e MAX_JOBS="$(nproc --ignore=2)" \ - -e SCCACHE_BUCKET \ - -e SCCACHE_REGION \ - -e SCCACHE_S3_KEY_PREFIX \ - -e XLA_CUDA \ - -e XLA_CLANG_CACHE_S3_BUCKET_NAME \ - -e PYTORCH_TEST_CUDA_MEM_LEAK_CHECK \ - -e PYTORCH_TEST_RERUN_DISABLED_TESTS \ - -e SKIP_SCCACHE_INITIALIZATION=1 \ - -e HUGGING_FACE_HUB_TOKEN \ - -e SCRIBE_GRAPHQL_ACCESS_TOKEN \ - -e DASHBOARD_TAG \ - -e IS_A100_RUNNER \ - --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \ - --security-opt seccomp=unconfined \ - --cap-add=SYS_PTRACE \ - --ipc=host \ - --shm-size="${SHM_SIZE}" \ - --tty \ - --detach \ - --name="${container_name}" \ - --user jenkins \ - -v "${GITHUB_WORKSPACE}:/var/lib/jenkins/workspace" \ - -w /var/lib/jenkins/workspace \ - "${DOCKER_IMAGE}" - ) - echo "DOCKER_CONTAINER_ID=${container_name}" >> "${GITHUB_ENV}" - docker exec -t "${container_name}" sh -c "pip install $(echo dist/*.whl)[opt-einsum] && ${TEST_COMMAND}" - - - name: Upload pytest cache if tests failed - uses: ./.github/actions/pytest-cache-upload - continue-on-error: true - if: failure() && steps.test.conclusion && steps.test.conclusion == 'failure' - with: - cache_dir: .pytest_cache - shard: ${{ matrix.shard }} - sha: ${{ github.event.pull_request.head.sha || github.sha }} - test_config: ${{ matrix.config }} - job_identifier: ${{ github.workflow }}_${{ inputs.build-environment }} - - - name: Print remaining test logs - shell: bash - if: always() && steps.test.conclusion - run: | - cat test/**/*_toprint.log || true - - - name: Stop monitoring script - if: ${{ always() && steps.monitor-script.outputs.monitor-script-pid }} - shell: bash - continue-on-error: true - env: - MONITOR_SCRIPT_PID: ${{ steps.monitor-script.outputs.monitor-script-pid }} - run: | - kill "$MONITOR_SCRIPT_PID" - - - name: Upload test artifacts - uses: ./.github/actions/upload-test-artifacts - if: always() && steps.test.conclusion && steps.test.conclusion != 'skipped' - with: - file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.get-job-id.outputs.job-id }} - use-gha: ${{ inputs.use-gha }} - s3-bucket: ${{ inputs.s3-bucket }} - - - name: Collect backtraces from coredumps (if any) - if: always() - shell: bash - run: | - # shellcheck disable=SC2156 - find . -iname "core.[1-9]*" -exec docker exec "${DOCKER_CONTAINER_ID}" sh -c "gdb python {} -ex 'bt' -ex 'q'" \; - - - name: Store Core dumps on S3 - uses: seemethere/upload-artifact-s3@v5 - if: failure() - with: - name: coredumps-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }} - retention-days: 14 - if-no-files-found: ignore - path: ./**/core.[1-9]* - - - name: Teardown Linux - uses: pytorch/test-infra/.github/actions/teardown-linux@main - if: always() && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' - - # NB: We are currently having an intermittent GPU-related issue on G5 runners with - # A10G GPU. Once this happens, trying to reset the GPU as done in setup-nvidia does - # not seem to help. Here are some symptoms: - # * Calling nvidia-smi timeouts after 60 second - # * Fail to run nvidia-smi with an unable to determine the device handle for GPU - # unknown error - # * Test fails with a missing CUDA GPU error when initializing CUDA in PyTorch - # * Run docker --gpus all fails with error response from daemon - # - # As both the root cause and recovery path are unclear, let's take the runner out of - # service so that it doesn't get any more jobs - - name: Check NVIDIA driver installation step - if: failure() && steps.install-nvidia-driver.outcome && steps.install-nvidia-driver.outcome != 'skipped' - shell: bash - env: - RUNNER_WORKSPACE: ${{ runner.workspace }} - run: | - set +e - set -x - - nvidia-smi - # NB: Surprisingly, nvidia-smi command returns successfully with return code 0 even in - # the case where the driver has already crashed as it still can get the driver version - # and some basic information like the bus ID. However, the rest of the information - # would be missing (ERR!), for example: - # - # +-----------------------------------------------------------------------------+ - # | NVIDIA-SMI 525.89.02 Driver Version: 525.89.02 CUDA Version: 12.0 | - # |-------------------------------+----------------------+----------------------+ - # | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | - # | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | - # | | | MIG M. | - # |===============================+======================+======================| - # | 0 ERR! Off | 00000000:00:1E.0 Off | ERR! | - # |ERR! ERR! ERR! ERR! / ERR! | 4184MiB / 23028MiB | ERR! Default | - # | | | ERR! | - # +-------------------------------+----------------------+----------------------+ - # - # +-----------------------------------------------------------------------------+ - # | Processes: | - # | GPU GI CI PID Type Process name GPU Memory | - # | ID ID Usage | - # |=============================================================================| - # +-----------------------------------------------------------------------------+ - # - # This should be reported as a failure instead as it will guarantee to fail when - # Docker tries to run with --gpus all - # - # So, the correct check here is to query one of the missing piece of info like - # GPU name, so that the command can fail accordingly - nvidia-smi --query-gpu=gpu_name --format=csv,noheader --id=0 - NVIDIA_SMI_STATUS=$? - - # These are acceptable return code from nvidia-smi as copied from setup-nvidia GitHub action - if [ "$NVIDIA_SMI_STATUS" -ne 0 ] && [ "$NVIDIA_SMI_STATUS" -ne 14 ]; then - echo "NVIDIA driver installation has failed, shutting down the runner..." - .github/scripts/stop_runner_service.sh - fi - - # For runner with multiple GPUs, we also want to confirm that the number of GPUs are the - # power of 2, i.e. 1, 2, 4, or 8. This is to avoid flaky test issue when one GPU fails - # https://github.com/pytorch/test-infra/issues/4000 - GPU_COUNT=$(nvidia-smi --list-gpus | wc -l) - NVIDIA_SMI_STATUS=$? - - # These are acceptable return code from nvidia-smi as copied from setup-nvidia GitHub action - if [ "$NVIDIA_SMI_STATUS" -ne 0 ] && [ "$NVIDIA_SMI_STATUS" -ne 14 ]; then - echo "NVIDIA driver installation has failed, shutting down the runner..." - .github/scripts/stop_runner_service.sh - fi - - # Check the GPU count to be a power of 2 - if [ "$GPU_COUNT" -le 8 ] && [ "$GPU_COUNT" -ne 1 ] && [ "$GPU_COUNT" -ne 2 ] && [ "$GPU_COUNT" -ne 4 ] && [ "$GPU_COUNT" -ne 8 ]; then - echo "NVIDIA driver detects $GPU_COUNT GPUs. The runner has a broken GPU, shutting it down..." - .github/scripts/stop_runner_service.sh - fi diff --git a/.github/actions/setup-linux/action.yml b/.github/actions/setup-linux/action.yml index 489a5b9828e5c..e9f0b74cbb97b 100644 --- a/.github/actions/setup-linux/action.yml +++ b/.github/actions/setup-linux/action.yml @@ -19,10 +19,79 @@ inputs: description: CUDA version to activate (e.g. "13.0"). Empty string keeps the default. required: false default: '' + submodules: + description: Submodule checkout mode passed to checkout-pytorch (default "recursive", use "false" for test jobs). + required: false + default: 'recursive' + github-token: + description: GITHUB_TOKEN, needed to retrieve the workflow job id. + required: false + default: '' + +outputs: + branch: + description: Parsed branch name from GITHUB_REF + value: ${{ steps.parse-ref.outputs.branch }} + tag: + description: Parsed tag name from GITHUB_REF (if applicable) + value: ${{ steps.parse-ref.outputs.tag }} + job-id: + description: The workflow job id + value: ${{ steps.get-job-id.outputs.job-id }} + job-name: + description: The workflow job name + value: ${{ steps.get-job-id.outputs.job-name }} runs: using: composite steps: + # ── ARC-only steps ────────────────────────────────────────────────── + - name: Fix workspace permissions + if: ${{ inputs.use-arc == 'true' }} + shell: bash + run: | + # GH runner image has switched to uid 1001 https://github.com/actions/runner-images/issues/10936 + # while current PyTorch CI image are still using uid 1000 (ec2-user). We + # can update the uid to 1001 eventually when everything migrates to ARC. + # In the meantime, this is a quick fix to ensure that CI has the permission + # to use the GITHUB_WORKSPACE while still allowing the GH hook (uid 1001, + # gid 1001) to clean up the directory after the job + sudo chmod -R 777 "$GITHUB_WORKSPACE" + + - name: Ack Git cache ownership + if: ${{ inputs.use-arc == 'true' }} + shell: bash + run: | + git config --global --add safe.directory "$GITHUB_WORKSPACE" + + # ── Common steps (shared) ──────────────────────────────────────── + - name: Checkout PyTorch + uses: pytorch/pytorch/.github/actions/checkout-pytorch@main + with: + no-sudo: true + checkout-mode: treeless + submodules: ${{ inputs.submodules }} + + - name: Parse ref + id: parse-ref + shell: bash + run: | + if [ -f "${{ github.action_path }}/../../scripts/parse_ref.py" ]; then + python3 "${{ github.action_path }}/../../scripts/parse_ref.py" + elif [ -f .github/scripts/parse_ref.py ]; then + python3 .github/scripts/parse_ref.py + else + echo "ERROR: parse_ref.py not found" >&2 + exit 1 + fi + + - name: Get workflow job id + id: get-job-id + if: ${{ always() && inputs.github-token != '' }} + uses: pytorch/pytorch/.github/actions/get-workflow-job-id@main + with: + github-token: ${{ inputs.github-token }} + # ── EC2-only steps ────────────────────────────────────────────────── - name: Display EC2 information if: ${{ inputs.use-arc != 'true' }} @@ -56,7 +125,7 @@ runs: - name: Start docker if docker daemon is not running shell: bash - if: ${{ inputs.use-arc != 'true' && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }} + if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }} run: | if ! docker version >/dev/null 2>/dev/null; then if systemctl is-active --quiet docker; then @@ -68,13 +137,13 @@ runs: - name: Install uv (EC2) uses: pytorch/test-infra/.github/actions/setup-uv@main - if: ${{ inputs.use-arc != 'true' && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }} + if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }} with: python-version: "3.12" - name: Install pip shell: bash - if: ${{ inputs.use-arc != 'true' && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }} + if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }} run: | set -euo pipefail if command -v pip3 &> /dev/null; then @@ -85,7 +154,7 @@ runs: fi - name: Kill any existing containers, clean up images - if: ${{ inputs.use-arc != 'true' && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }} + if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }} shell: bash run: | # ignore expansion of "docker ps -q" since it could be empty @@ -97,7 +166,7 @@ runs: - name: Check that the docker daemon is running shell: bash continue-on-error: true - if: ${{ inputs.use-arc != 'true' && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'true' }} + if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'true' }} run: | set +x @@ -119,113 +188,6 @@ runs: echo "Reached maximum attempts to connect to Docker. Exiting." exit 1 - # ── ARC-only steps ────────────────────────────────────────────────── - - name: Install uv (ARC) - if: ${{ inputs.use-arc == 'true' }} - uses: pytorch/test-infra/.github/actions/setup-uv@main - with: - python-version: ${{ inputs.python-version || '3.12' }} - activate-environment: true - - - name: Install CI requirements - if: ${{ inputs.use-arc == 'true' }} - shell: bash - run: | - uv pip install -r .ci/docker/requirements-ci.txt - echo "CMAKE_PREFIX_PATH=${VIRTUAL_ENV}" >> "${GITHUB_ENV}" - - - name: Install C++ compiler - if: ${{ inputs.use-arc == 'true' && inputs.compiler != '' }} - shell: bash - run: | - set -eux - - COMPILER="${{ inputs.compiler }}" - CLANG_VERSION=$(echo "${COMPILER}" | perl -n -e'/^clang(\d+)$/ && print $1') - GCC_VERSION=$(echo "${COMPILER}" | perl -n -e'/^gcc(\d+)$/ && print $1') - - if [ -z "${CLANG_VERSION}" ] && [ -z "${GCC_VERSION}" ]; then - echo "ERROR: unrecognized compiler '${COMPILER}', expected gcc or clang" - exit 1 - elif [ -n "${CLANG_VERSION}" ]; then - # Clang 20 and libomp-devel 20 are already installed in the base image. - # For older versions (15, 18), download pre-built tarballs from S3. - if [ "${CLANG_VERSION}" = "20" ]; then - echo "Using base image default clang 20" - elif [ "${CLANG_VERSION}" = "15" ] || [ "${CLANG_VERSION}" = "18" ]; then - ARCH=$(uname -m) - PLATFORM="linux-${ARCH}" - S3_URL="https://gha-artifacts.s3.amazonaws.com/clang/${PLATFORM}/clang-${CLANG_VERSION}/clang-${CLANG_VERSION}-${PLATFORM}.tar.gz" - INSTALL_DIR="/opt/clang-${CLANG_VERSION}" - - echo "Downloading clang ${CLANG_VERSION} from ${S3_URL}" - curl -fsSL "${S3_URL}" -o /tmp/clang.tar.gz - mkdir -p "${INSTALL_DIR}" - tar xzf /tmp/clang.tar.gz -C "${INSTALL_DIR}" --strip-components=1 - rm -f /tmp/clang.tar.gz - - echo "PATH=${INSTALL_DIR}/bin:${PATH}" >> "${GITHUB_ENV}" - echo "CC=${INSTALL_DIR}/bin/clang" >> "${GITHUB_ENV}" - echo "CXX=${INSTALL_DIR}/bin/clang++" >> "${GITHUB_ENV}" - echo "CMAKE_PREFIX_PATH=${INSTALL_DIR}:${CMAKE_PREFIX_PATH:-}" >> "${GITHUB_ENV}" - else - echo "ERROR: unsupported clang version '${CLANG_VERSION}', available: 15, 18, 20" - exit 1 - fi - else - # On AlmaLinux 9, GCC 11 is the system default at /usr/bin/gcc. - # Only install gcc-toolset for newer versions that aren't already available. - if [ "${GCC_VERSION}" = "11" ]; then - # Remove gcc/g++ from /usr/local/bin so they don't shadow /usr/bin/gcc (system default) - rm -f /usr/local/bin/gcc /usr/local/bin/g++ /usr/local/bin/cc /usr/local/bin/c++ - # Strip any /opt/rh/gcc-toolset-* paths from PATH - CLEAN_PATH=$(echo "${PATH}" | tr ':' '\n' | grep -v '/opt/rh/gcc-toolset-' | tr '\n' ':' | sed 's/:$//') - echo "PATH=${CLEAN_PATH}" >> "${GITHUB_ENV}" - echo "Using system default GCC 11" - else - dnf install -y \ - gcc-toolset-${GCC_VERSION}-gcc \ - gcc-toolset-${GCC_VERSION}-gcc-c++ \ - gcc-toolset-${GCC_VERSION}-gcc-gfortran - - TOOLSET_ROOT="/opt/rh/gcc-toolset-${GCC_VERSION}/root/usr" - echo "PATH=${TOOLSET_ROOT}/bin:${PATH}" >> "${GITHUB_ENV}" - echo "LD_LIBRARY_PATH=${TOOLSET_ROOT}/lib64:${TOOLSET_ROOT}/lib:${LD_LIBRARY_PATH}" >> "${GITHUB_ENV}" - fi - fi - - - name: Switch CUDA version - if: ${{ inputs.use-arc == 'true' && inputs.cuda-version != '' }} - shell: bash - run: | - set -eux - - CUDA_VERSION="${{ inputs.cuda-version }}" - - # All CUDA versions are pre-installed in the base image under - # /usr/local/cuda-{major.minor}. The default (/usr/local/cuda - # symlink) points to CUDA 13. Switch to the requested version. - CUDA_DIR="/usr/local/cuda-${CUDA_VERSION}" - if [ ! -d "${CUDA_DIR}" ]; then - echo "ERROR: CUDA ${CUDA_VERSION} not found at ${CUDA_DIR}" - ls -d /usr/local/cuda-* /usr/local/cuda 2>/dev/null || true - exit 1 - fi - - rm -f /usr/local/cuda - ln -s "${CUDA_DIR}" /usr/local/cuda - - echo "PATH=/usr/local/cuda/bin:${PATH}" >> "${GITHUB_ENV}" - echo "LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH:-}" >> "${GITHUB_ENV}" - echo "CUDA_HOME=/usr/local/cuda" >> "${GITHUB_ENV}" - echo "CUDA_PATH=/usr/local/cuda" >> "${GITHUB_ENV}" - echo "DESIRED_CUDA=${CUDA_VERSION}" >> "${GITHUB_ENV}" - echo "NCCL_INCLUDE_DIR=/usr/local/cuda/include" >> "${GITHUB_ENV}" - echo "NCCL_LIB_DIR=/usr/local/cuda/lib64" >> "${GITHUB_ENV}" - - echo "Switched to CUDA ${CUDA_VERSION}" - nvcc --version - # ── Shared steps ──────────────────────────────────────────────────── - name: Print GPU info (if present) shell: bash diff --git a/.github/actions/setup-rocm/action.yml b/.github/actions/setup-rocm/action.yml index 5bb982a4085b1..3c1445b85fa84 100644 --- a/.github/actions/setup-rocm/action.yml +++ b/.github/actions/setup-rocm/action.yml @@ -110,7 +110,7 @@ runs: # This is due to the device files (/dev/kfd & /dev/dri) being owned by video group on bare metal. # This video group ID maps to subgid 1 inside the docker image due to the /etc/subgid entries. # The group name corresponding to group ID 1 can change depending on the OS, so both are necessary. - echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd $DEVICE_FLAG --group-add video --group-add $render_gid --group-add daemon --group-add bin --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --network=host" >> "${GITHUB_ENV}" + echo "GPU_FLAG=--device=/dev/mem --device=/dev/kfd $DEVICE_FLAG --group-add video --group-add $render_gid --group-add daemon --group-add bin --cap-add=SYS_PTRACE --security-opt seccomp=unconfined" >> "${GITHUB_ENV}" - name: Login to ECR uses: pytorch/pytorch/.github/actions/ecr-login@main diff --git a/.github/actions/setup-xpu/action.yml b/.github/actions/setup-xpu/action.yml index dfae0a358259e..23b750205ab08 100644 --- a/.github/actions/setup-xpu/action.yml +++ b/.github/actions/setup-xpu/action.yml @@ -44,7 +44,7 @@ runs: fi - name: Runner diskspace health check - uses: ./.github/actions/diskspace-cleanup + uses: pytorch/pytorch/.github/actions/diskspace-cleanup@main if: always() - name: Runner health check disconnect on failure @@ -53,11 +53,30 @@ runs: run: | killall runsvc.sh + - name: Setup useful environment variables + shell: bash + run: | + RUNNER_ARTIFACT_DIR="${RUNNER_TEMP}/artifacts" + rm -rf "${RUNNER_ARTIFACT_DIR}" + mkdir -p "${RUNNER_ARTIFACT_DIR}" + echo "RUNNER_ARTIFACT_DIR=${RUNNER_ARTIFACT_DIR}" >> "${GITHUB_ENV}" + + RUNNER_TEST_RESULTS_DIR="${RUNNER_TEMP}/test-results" + rm -rf "${RUNNER_TEST_RESULTS_DIR}" + mkdir -p "${RUNNER_TEST_RESULTS_DIR}" + echo "RUNNER_TEST_RESULTS_DIR=${RUNNER_TEST_RESULTS_DIR}" >> "${GITHUB_ENV}" + + RUNNER_DOCS_DIR="${RUNNER_TEMP}/docs" + rm -rf "${RUNNER_DOCS_DIR}" + mkdir -p "${RUNNER_DOCS_DIR}" + echo "RUNNER_DOCS_DIR=${RUNNER_DOCS_DIR}" >> "${GITHUB_ENV}" + - name: Preserve github env variables for use in docker shell: bash run: | - env | grep '^GITHUB' >> "/tmp/github_env_${GITHUB_RUN_ID}" - env | grep '^CI' >> "/tmp/github_env_${GITHUB_RUN_ID}" + env | grep '^GITHUB' >> "${RUNNER_TEMP}/github_env_${GITHUB_RUN_ID}" + env | grep '^CI' >> "${RUNNER_TEMP}/github_env_${GITHUB_RUN_ID}" + env | grep '^RUNNER' >> "${RUNNER_TEMP}/github_env_${GITHUB_RUN_ID}" - name: XPU set GPU_FLAG shell: bash @@ -67,4 +86,4 @@ runs: echo "GPU_FLAG=--device=/dev/mem --device=/dev/dri --group-add video --group-add $render_gid" >> "${GITHUB_ENV}" - name: Login to ECR - uses: ./.github/actions/ecr-login + uses: pytorch/pytorch/.github/actions/ecr-login@main diff --git a/.github/actions/upload-build-artifacts/action.yml b/.github/actions/upload-build-artifacts/action.yml new file mode 100644 index 0000000000000..62413131581ea --- /dev/null +++ b/.github/actions/upload-build-artifacts/action.yml @@ -0,0 +1,37 @@ +name: Upload PyTorch Build Artifacts + +description: Upload build artifacts to S3 or GHA. + +inputs: + name: + description: Name of the artifact + required: true + use-gha: + description: If set to any value, use GHA to upload the artifact. Otherwise use S3. + required: false + s3-bucket: + description: S3 bucket to upload builds + required: false + default: "gha-artifacts" + +runs: + using: composite + steps: + - name: Upload PyTorch Build Artifacts to S3 + if: inputs.use-gha == '' || inputs.use-gha == 'false' + uses: seemethere/upload-artifact-s3@baba72d0712b404f646cebe0730933554ebce96a # v5.1.0 + with: + name: ${{ inputs.name }} + retention-days: 14 + if-no-files-found: error + path: artifacts.zip + s3-bucket: ${{ inputs.s3-bucket }} + + - name: Upload PyTorch Build Artifacts to GHA + if: inputs.use-gha != '' && inputs.use-gha != 'false' + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + with: + name: ${{ inputs.name }} + retention-days: 14 + if-no-files-found: error + path: artifacts.zip diff --git a/.github/actions/upload-test-artifacts/action.yml b/.github/actions/upload-test-artifacts/action.yml index 4029ae79441f1..9fcb125a42313 100644 --- a/.github/actions/upload-test-artifacts/action.yml +++ b/.github/actions/upload-test-artifacts/action.yml @@ -37,7 +37,7 @@ runs: FILE_SUFFIX: ${{ inputs.file-suffix }} run: | set -euo pipefail - rm -f test-jsons-*.zip test-reports-*.zip logs-*.zip debug-*.zip + rm -f test-jsons-*.zip test-reports-*.zip logs-*.zip debug-*.zip profiler-traces-*.zip tlparse-*.zip ZIP_CMD=' import sys, zipfile, os from pathlib import Path @@ -60,6 +60,14 @@ runs: uv run --no-project python -c "$ZIP_CMD" \ test/debug "debug-${FILE_SUFFIX}.zip" '**/*' fi + if [ -d 'test/test-reports/profiler_traces' ]; then + uv run --no-project python -c "$ZIP_CMD" \ + test/test-reports/profiler_traces "profiler-traces-${FILE_SUFFIX}.zip" '**/*' + fi + if [ -d 'test/test-reports/tlparse_output' ]; then + uv run --no-project python -c "$ZIP_CMD" \ + test/test-reports/tlparse_output "tlparse-${FILE_SUFFIX}.zip" '**/*' + fi # Windows zip - name: Zip JSONs for upload @@ -141,6 +149,30 @@ runs: if-no-files-found: ignore path: debug-*.zip + - name: Store Profiler Traces on S3 + id: s3-upload-profiler-traces + uses: seemethere/upload-artifact-s3@v5 + continue-on-error: true + with: + s3-bucket: ${{ inputs.s3-bucket }} + s3-prefix: | + ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact + retention-days: 14 + if-no-files-found: ignore + path: profiler-traces-*.zip + + - name: Store TLParse Output on S3 + id: s3-upload-tlparse + uses: seemethere/upload-artifact-s3@v5 + continue-on-error: true + with: + s3-bucket: ${{ inputs.s3-bucket }} + s3-prefix: | + ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact + retention-days: 14 + if-no-files-found: ignore + path: tlparse-*.zip + # Check if S3 upload failed (test-reports is the critical one) - name: Check S3 upload status id: check-s3 @@ -193,3 +225,21 @@ runs: path: | usage_log.txt test/**/*.log + + - name: Store Profiler Traces on Github + uses: actions/upload-artifact@v4 + continue-on-error: true + with: + name: profiler-traces-runattempt${{ github.run_attempt }}-${{ inputs.file-suffix }} + retention-days: 14 + if-no-files-found: ignore + path: test/test-reports/profiler_traces/**/* + + - name: Store TLParse Output on Github + uses: actions/upload-artifact@v4 + continue-on-error: true + with: + name: tlparse-runattempt${{ github.run_attempt }}-${{ inputs.file-suffix }} + retention-days: 14 + if-no-files-found: ignore + path: test/test-reports/tlparse_output/**/* diff --git a/.github/allowlist.yml b/.github/allowlist.yml new file mode 100644 index 0000000000000..c6a0bd196837c --- /dev/null +++ b/.github/allowlist.yml @@ -0,0 +1,36 @@ +# Cross Repo CI Relay (CRCR) Allowlist +# +# This document defines which downstream repositories can receive PyTorch +# PR events or feedback through CRCR. Each repository is assigned a level +# (L1-L4), which determines the depth to which downstream repositories can +# participate. +# +# Trust Levels: +# +# L1 – Onboarding +# Events are forwarded to downstream, but upstream receives no feedback. +# +# L2 – Observation +# Downstream CI results are displayed on the HUD page, but not on PRs. +# +# L3 – Stable +# Adds a non-blocking check run on PRs when ciflow/oot/ label is applied. +# +# L4 – Mature +# Adds a blocking check run on every PR; reserved for critical accelerators. +# +# For more information, see the RFC: https://github.com/pytorch/rfcs/pull/90 + +# Example: +# +# L1: +# - org1/downstream-repo1 +# +# L2: +# - org2/downstream-repo2 +# +# L3: +# - org3/downstream-repo3: @oncall1,oncall2 +# +# L4: +# - org4/downstream-repo4: @oncall1,oncall2 diff --git a/.github/arc.yaml b/.github/arc.yaml index 0a74e2826ab0b..497f4f5f01e9c 100644 --- a/.github/arc.yaml +++ b/.github/arc.yaml @@ -25,85 +25,65 @@ runner_mapping: # ---- x86 CPU — Intel AVX-512 (c5, c7i families) ---- linux.large: l-x86iavx512-2-4 # c5.large - linux.c7i.large: l-x86iavx512-2-4 # c7i.large - linux.2xlarge: l-x86iavx512-8-16 # c5.2xlarge - linux.c7i.2xlarge: l-x86iavx512-8-16 # c7i.2xlarge - linux.4xlarge: l-x86iavx512-16-32 # c5.4xlarge - linux.4xlarge.for.testing.donotuse: l-x86iavx512-16-32 # c5.4xlarge - linux.c7i.4xlarge: l-x86iavx512-16-32 # c7i.4xlarge - linux.c7i.8xlarge: l-x86iavx512-48-96 # c7i.8xlarge — upgraded (no 32-64 equivalent) - linux.9xlarge.ephemeral: l-x86iavx512-36-72 # c5.9xlarge - linux.12xlarge: l-x86iavx512-48-96 # c5.12xlarge - linux.12xlarge.ephemeral: l-x86iavx512-48-96 # c5.12xlarge - linux.c7i.12xlarge: l-x86iavx512-48-96 # c7i.12xlarge - linux.16xlarge.spr: l-x86iavx512-94-192 # c7i.16xlarge — upgraded (no 64-128 equivalent) - linux.24xlarge: l-x86iavx512-94-192 # c5.24xlarge - linux.24xlarge.ephemeral: l-x86iavx512-94-192 # c5.24xlarge - linux.c7i.24xlarge: l-x86iavx512-94-192 # c7i.24xlarge - linux.24xl.spr-metal: l-bx86iamx-94-192 # c7i.metal-24xl + linux.2xlarge: l-x86iavx512-8-64 # c5.2xlarge + linux.c7i.2xlarge: l-x86iavx512-8-64 # c7i.2xlarge + linux.4xlarge: l-x86iavx512-16-128 # c5.4xlarge + linux.c7i.4xlarge: l-x86iavx512-16-128 # c7i.4xlarge + linux.12xlarge: l-x86iavx512-48-384 # c5.12xlarge + linux.c7i.12xlarge: l-x86iavx512-48-384 # c7i.12xlarge + linux.24xl.spr-metal: l-bx86iamx-92-167 # c7i.metal-24xl # ---- x86 CPU — Intel AMX (m7i-flex family) ---- - linux.2xlarge.amx: l-x86iamx-8-32 # m7i-flex.2xlarge - linux.4xlarge.amx: l-x86iamx-32-128 # m7i-flex.4xlarge — upgraded (no 16-64 equivalent) + linux.2xlarge.amx: l-x86iamx-8-64 # m7i-flex.2xlarge linux.8xlarge.amx: l-x86iamx-32-128 # m7i-flex.8xlarge # ---- x86 CPU — Intel AVX2 (m4 family) ---- linux.2xlarge.avx2: l-x86iavx2-8-32 # m4.2xlarge - linux.4xlarge.avx2: l-x86iavx2-40-160 # m4.4xlarge — upgraded (no 16-64 equivalent) linux.10xlarge.avx2: l-x86iavx2-40-160 # m4.10xlarge # ---- x86 CPU — Memory-optimized (r5, r7i families) ---- - linux.r7i.large: l-x86iavx512-8-64 # r7i.large — upgraded (no 2-16 equivalent) - linux.r7i.xlarge: l-x86iavx512-8-64 # r7i.xlarge — upgraded (no 4-32 equivalent) linux.r7i.2xlarge: l-x86iavx512-8-64 # r7i.2xlarge linux.r7i.4xlarge: l-x86iavx512-16-128 # r7i.4xlarge - linux.r7i.8xlarge: l-x86iavx512-32-256 # r7i.8xlarge - linux.r7i.12xlarge: l-x86iavx512-48-384 # r7i.12xlarge - linux.2xlarge.memory: l-x86iavx512-8-64 # r5.2xlarge linux.4xlarge.memory: l-x86iavx512-16-128 # r5.4xlarge linux.8xlarge.memory: l-x86iavx512-32-256 # r5.8xlarge linux.12xlarge.memory: l-x86iavx512-48-384 # r5.12xlarge - linux.12xlarge.memory.ephemeral: l-x86iavx512-48-384 # r5.12xlarge - linux.16xlarge.memory: l-x86iavx512-94-768 # r5.16xlarge — upgraded (no 64-512 equivalent) linux.24xlarge.memory: l-x86iavx512-94-768 # r5.24xlarge # ---- x86 CPU — AMD (m6a, m7a families) ---- - linux.8xlarge.amd: l-x86aavx512-94-384 # m7a.8xlarge — upgraded (no 32-128 equivalent) - linux.12xlarge.amd: l-x86aavx512-94-384 # m6a.12xlarge — upgraded (no 48-192 equivalent) - linux.24xlarge.amd: l-x86aavx512-94-384 # m7a.24xlarge + linux.24xlarge.amd: l-x86aavx512-125-463 # m7a.24xlarge # ---- x86 GPU — T4 (g4dn family) ---- - linux.g4dn.4xlarge.nvidia.gpu: l-x86iavx512-16-64-t4 # g4dn.4xlarge - linux.g4dn.12xlarge.nvidia.gpu: l-x86iavx512-48-192-t4-4 # g4dn.12xlarge - linux.g4dn.metal.nvidia.gpu: l-bx86iavx512-94-384-t4-8 # g4dn.metal + linux.g4dn.4xlarge.nvidia.gpu: l-x86iavx512-29-115-t4 # g4dn.4xlarge + linux.g4dn.12xlarge.nvidia.gpu: l-x86iavx512-45-172-t4-4 # g4dn.12xlarge + linux.g4dn.metal.nvidia.gpu: l-bx86iavx512-94-344-t4-8 # g4dn.metal # ---- x86 GPU — A10G (g5 family) ---- - linux.g5.4xlarge.nvidia.gpu: l-x86aavx2-16-64-a10g # g5.4xlarge - linux.g5.12xlarge.nvidia.gpu: l-x86aavx2-48-192-a10g-4 # g5.12xlarge - linux.g5.48xlarge.nvidia.gpu: l-x86aavx2-192-768-a10g-8 # g5.48xlarge + linux.g5.4xlarge.nvidia.gpu: l-x86aavx2-29-113-a10g # g5.4xlarge + linux.g5.12xlarge.nvidia.gpu: l-x86aavx2-45-167-a10g-4 # g5.12xlarge + linux.g5.48xlarge.nvidia.gpu: l-x86aavx2-189-704-a10g-8 # g5.48xlarge # ---- x86 GPU — L4 (g6 family) ---- - linux.g6.4xlarge.experimental.nvidia.gpu: l-x86aavx2-16-64-l4 # g6.4xlarge - linux.g6.12xlarge.nvidia.gpu: l-x86aavx2-48-192-l4-4 # g6.12xlarge - - # ---- x86 GPU — V100 (p3 family) ---- - - linux.p3.8xlarge.nvidia.gpu: l-x86aavx2-48-192-a10g-4 # p3.8xlarge — upgraded (no V100 equivalent; 4x A10G closest match) + linux.g6.4xlarge.experimental.nvidia.gpu: l-x86aavx2-29-113-l4 # g6.4xlarge + linux.g6.12xlarge.nvidia.gpu: l-x86aavx2-45-172-l4-4 # g6.12xlarge # ---- ARM64 — Graviton ---- linux.arm64.2xlarge: l-arm64g2-6-32 # t4g.2xlarge linux.arm64.2xlarge.ephemeral: l-arm64g2-6-32 # t4g.2xlarge - linux.arm64.m7g.4xlarge: l-arm64g3-16-64 # m7g.4xlarge - linux.arm64.m7g.4xlarge.ephemeral: l-arm64g3-16-64 # m7g.4xlarge - linux.arm64.m8g.4xlarge: l-arm64g4-16-64 # m8g.4xlarge - linux.arm64.m8g.4xlarge.ephemeral: l-arm64g4-16-64 # m8g.4xlarge - linux.arm64.r7g.12xlarge.memory: l-arm64g3-48-384 # r7g.12xlarge - linux.arm64.m7g.metal: l-barm64g3-62-256 # m7g.metal + linux.arm64.m7g.4xlarge: l-arm64g3-16-62 # m7g.4xlarge + linux.arm64.m8g.4xlarge: l-arm64g4-16-62 # m8g.4xlarge + linux.arm64.r7g.12xlarge.memory: l-arm64g3-61-463 # r7g.12xlarge + linux.arm64.m7g.metal: l-barm64g4-62-226 # m7g.metal + + + # ---- Partner hardwares ---- + + linux.idc.xpu: linux.idc.xpu + linux.rocm.gpu.2: linux.rocm.gpu.2 diff --git a/.github/ci_commit_pins/audio.txt b/.github/ci_commit_pins/audio.txt index ad631f37d819f..e5487b7ac5eff 100644 --- a/.github/ci_commit_pins/audio.txt +++ b/.github/ci_commit_pins/audio.txt @@ -1 +1 @@ -11ed3578e70404c7329f219b02d78b7a89603ebe +c0cbdb95674556cdff7266f2d44bb855f634cfde diff --git a/.github/ci_commit_pins/torchtitan.txt b/.github/ci_commit_pins/torchtitan.txt index f086735e52eb1..0c66dcaabc888 100644 --- a/.github/ci_commit_pins/torchtitan.txt +++ b/.github/ci_commit_pins/torchtitan.txt @@ -1 +1 @@ -a079987c2974857a4fb5739f7270c8392d257815 +7c8fc3f540edaec84d4610066ccff423141f7aa7 diff --git a/.github/ci_commit_pins/vision.txt b/.github/ci_commit_pins/vision.txt index 94cddf66c7ee9..516c2c1931b15 100644 --- a/.github/ci_commit_pins/vision.txt +++ b/.github/ci_commit_pins/vision.txt @@ -1 +1 @@ -a025873a304c940c5117ef38e8bb778c55a06f7f +601776dc14ab12179412fa7fb08762e20862720c diff --git a/.github/ci_commit_pins/vllm.txt b/.github/ci_commit_pins/vllm.txt index 3db4b7d5c4300..21f23c297a01d 100644 --- a/.github/ci_commit_pins/vllm.txt +++ b/.github/ci_commit_pins/vllm.txt @@ -1 +1 @@ -5db91f0aaf3566b1d9f8b0720065eb9009296d98 +ba4a78eb5d2ea30477b58a0bb8109b129f35c8b1 diff --git a/.github/labeler.yml b/.github/labeler.yml index f1d10d3a8082b..c8f77e91be65d 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -62,16 +62,6 @@ - test/test_jit_llga_fuser.py - test/test_mkldnn.py -"ciflow/linux-aarch64": -- third_party/ideep -- caffe2/ideep/** -- caffe2/python/ideep/** -- cmake/Modules/FindMKLDNN.cmake -- third_party/mkl-dnn.BUILD -- torch/csrc/jit/codegen/onednn/** -- test/test_jit_llga_fuser.py -- test/test_mkldnn.py - "module: amp (automated mixed precision)": - torch/amp/** - aten/src/ATen/autocast_mode.* @@ -200,6 +190,8 @@ - test/distributed/**/*mem*/** "ciflow/torchtitan": +# torchtitan commit pin updates +- .github/ci_commit_pins/torchtitan.txt # torch.distributed (FSDP, DTensor, etc.) - torch/distributed/** # torch.compile diff --git a/.github/pytorch-probot.yml b/.github/pytorch-probot.yml index f47c3bc005cd2..6f875dcf0716d 100644 --- a/.github/pytorch-probot.yml +++ b/.github/pytorch-probot.yml @@ -27,7 +27,6 @@ ciflow_push_tags: - ciflow/inductor-rocm-mi200 - ciflow/inductor-rocm-mi300 - ciflow/inductor-rocm-mi355 -- ciflow/linux-aarch64 - ciflow/mps - ciflow/nightly - ciflow/op-benchmark @@ -60,6 +59,8 @@ retryable_workflows: - linux-binary - windows-binary - inductor-A100-perf-nightly +retryable_step_names: +- Initialize containers labeler_config: labeler.yml label_to_label_config: label_to_label.yml mergebot: true diff --git a/.github/requirements-gha-cache.txt b/.github/requirements-gha-cache.txt index c274ca1e5914d..9bd898687ba67 100644 --- a/.github/requirements-gha-cache.txt +++ b/.github/requirements-gha-cache.txt @@ -11,5 +11,5 @@ lintrunner==0.12.7 ninja==1.10.0.post1 nvidia-ml-py==11.525.84 pyyaml==6.0.2 -requests==2.32.4 +requests==2.33.0 rich==14.1.0 diff --git a/.github/scripts/amd/package_triton_wheel.sh b/.github/scripts/amd/package_triton_wheel.sh index 501e50e2fe2f1..549bacc5a7573 100755 --- a/.github/scripts/amd/package_triton_wheel.sh +++ b/.github/scripts/amd/package_triton_wheel.sh @@ -28,9 +28,7 @@ if [[ -z "${TRITON_ROCM_DIR}" ]]; then export TRITON_ROCM_DIR=third_party/amd/backend fi -# Remove packaged libs and headers -rm -rf $TRITON_ROCM_DIR/include/* - +# Remove packaged libs LIBNUMA_PATH="/usr/lib64/libnuma.so.1" LIBELF_PATH="/usr/lib64/libelf.so.1" OS_NAME=`awk -F= '/^NAME/{print $2}' /etc/os-release` diff --git a/.github/scripts/build_triton_wheel.py b/.github/scripts/build_triton_wheel.py index e5ac9c5937dfa..de0fa7739edc5 100644 --- a/.github/scripts/build_triton_wheel.py +++ b/.github/scripts/build_triton_wheel.py @@ -81,15 +81,23 @@ def build_triton( check_call(["git", "clone", triton_repo, "triton"], cwd=tmpdir) if release: ver, rev, patch = version.split(".") - check_call( - ["git", "checkout", f"release/{ver}.{rev}.x"], cwd=triton_basedir - ) + if device == "xpu": + # XPU uses the patch version in the release branch name + check_call( + ["git", "checkout", f"release/{ver}.{rev}.{patch}"], + cwd=triton_basedir, + ) + else: + check_call( + ["git", "checkout", f"release/{ver}.{rev}.x"], cwd=triton_basedir + ) else: check_call(["git", "fetch", "origin", commit_hash], cwd=triton_basedir) check_call(["git", "checkout", commit_hash], cwd=triton_basedir) # change built wheel name and version env["TRITON_WHEEL_NAME"] = triton_pkg_name + env["TRITON_EXT_ENABLED"] = "ON" if with_clang_ldd: env["TRITON_BUILD_WITH_CLANG_LLD"] = "1" diff --git a/.github/scripts/filter_test_configs.py b/.github/scripts/filter_test_configs.py index 5393f50673ef7..9f7f7a3a7bedc 100755 --- a/.github/scripts/filter_test_configs.py +++ b/.github/scripts/filter_test_configs.py @@ -30,19 +30,31 @@ logging.basicConfig(level=logging.INFO) -def is_cuda_or_rocm_job(job_name: str | None) -> bool: - if not job_name: - return False +def is_cuda_or_rocm_job( + job_name: str | None, config: dict[str, Any] | None = None +) -> bool: + if job_name and ("cuda" in job_name or "rocm" in job_name): + return True - return "cuda" in job_name or "rocm" in job_name + # Also check the runner name in the config, since some workflows (e.g. + # inductor-unittest) use job names that don't include "cuda" even though + # they target CUDA runners. + if config: + runner = config.get("runner", "") + if "nvidia.gpu" in runner or "rocm.gpu" in runner: + return True + + return False # Supported modes when running periodically. Only applying the mode when -# its lambda condition returns true -SUPPORTED_PERIODICAL_MODES: dict[str, Callable[[str | None], bool]] = { +# its lambda condition returns true. Each callable receives (job_name, config). +SUPPORTED_PERIODICAL_MODES: dict[ + str, Callable[[str | None, dict[str, Any] | None], bool] +] = { # Memory leak check is only needed for CUDA and ROCm jobs which utilize GPU memory "mem_leak_check": is_cuda_or_rocm_job, - "rerun_disabled_tests": lambda job_name: True, + "rerun_disabled_tests": lambda job_name, config=None: True, } # The link to the published list of disabled jobs @@ -225,7 +237,7 @@ def set_periodic_modes( for config in test_matrix.get("include", []): for mode, cond in SUPPORTED_PERIODICAL_MODES.items(): - if not cond(job_name): + if not cond(job_name, config): continue cfg = config.copy() diff --git a/.github/scripts/generate_binary_build_matrix.py b/.github/scripts/generate_binary_build_matrix.py index 9511d6f6976d2..857d95ab33ea2 100644 --- a/.github/scripts/generate_binary_build_matrix.py +++ b/.github/scripts/generate_binary_build_matrix.py @@ -23,19 +23,15 @@ REPO_ROOT = SCRIPT_DIR.parent.parent -CUDA_ARCHES = ["12.6", "12.8", "12.9", "13.0", "13.2"] +CUDA_ARCHES = ["12.6", "13.0", "13.2"] CUDA_STABLE = "13.0" CUDA_ARCHES_FULL_VERSION = { "12.6": "12.6.3", - "12.8": "12.8.1", - "12.9": "12.9.1", "13.0": "13.0.2", "13.2": "13.2.0", } CUDA_ARCHES_CUDNN_VERSION = { "12.6": "9", - "12.8": "9", - "12.9": "9", "13.0": "9", "13.2": "9", } @@ -50,13 +46,14 @@ CUDA_AARCH64_ARCHES = [ "12.6-aarch64", - "12.8-aarch64", - "12.9-aarch64", "13.0-aarch64", "13.2-aarch64", ] +# WARNING: For CUDA 13.0, cublas is pinned to a version range rather +# than an exact version. A broken cublas release within that range will be +# silently pulled in. PYTORCH_EXTRA_INSTALL_REQUIREMENTS = { "12.6": ( "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | " # noqa: B950 @@ -66,24 +63,9 @@ "nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | " "nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'" ), - "12.8": ( - "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | " # noqa: B950 - "cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | " - "nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | " - "nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | " - "nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | " - "nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'" - ), - "12.9": ( - "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | " # noqa: B950 - "cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | " - "nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | " - "nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | " - "nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | " - "nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux'" - ), "13.0": ( - "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | " # noqa: B950 + "cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | " # noqa: B950 + "nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | " "cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | " "nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | " "nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | " @@ -316,12 +298,7 @@ def generate_libtorch_matrix( if arches is None: arches = ["cpu"] if os == "windows": - # TODO (huydhn): Only build CUDA 12.9 for Linux. This logic is to be cleaned up - # in 2.10 - windows_cuda_arches = CUDA_ARCHES.copy() - windows_cuda_arches.remove("12.9") - windows_cuda_arches.remove("13.2") - arches += windows_cuda_arches + arches += CUDA_ARCHES if libtorch_variants is None: libtorch_variants = [ "shared-with-deps", @@ -374,12 +351,7 @@ def generate_wheels_matrix( if os == "linux": arches += CUDA_ARCHES + ROCM_ARCHES + XPU_ARCHES elif os == "windows": - # TODO (huydhn): Only build CUDA 12.9 for Linux. This logic is to be cleaned up - # in 2.10 - windows_cuda_arches = CUDA_ARCHES.copy() - windows_cuda_arches.remove("12.9") - windows_cuda_arches.remove("13.2") - arches += windows_cuda_arches + XPU_ARCHES + arches += CUDA_ARCHES + XPU_ARCHES elif os == "linux-aarch64": # Separate new if as the CPU type is different and # uses different build/test scripts @@ -415,7 +387,7 @@ def generate_wheels_matrix( # cuda linux wheels require PYTORCH_EXTRA_INSTALL_REQUIREMENTS to install if ( - arch_version in ["13.2", "13.0", "12.9", "12.8", "12.6"] + arch_version in ["13.2", "13.0", "12.6"] and os == "linux" or arch_version in CUDA_AARCH64_ARCHES ): @@ -470,6 +442,9 @@ def generate_wheels_matrix( "pytorch_extra_install_requirements": ( PYTORCH_EXTRA_INSTALL_REQUIREMENTS["xpu"] if gpu_arch_type == "xpu" + else PYTORCH_EXTRA_INSTALL_REQUIREMENTS[CUDA_STABLE] + if gpu_arch_type == "cpu" + and os in ("windows", "macos-arm64") else "" ), } diff --git a/.github/scripts/generate_ci_workflows.py b/.github/scripts/generate_ci_workflows.py index 2dfe570dad339..b1e33b6038ec7 100755 --- a/.github/scripts/generate_ci_workflows.py +++ b/.github/scripts/generate_ci_workflows.py @@ -211,7 +211,7 @@ class OperatingSystem: os=OperatingSystem.MACOS_ARM64, package_type="wheel", build_configs=_MACOS_ARM64_WHEEL_CONFIGS, - macos_runner="macos-14-xlarge", + macos_runner="macos-26-xlarge", ciflow_config=CIFlowConfig( labels={ LABEL_CIFLOW_BINARIES, diff --git a/.github/scripts/map_ec2_to_arc.py b/.github/scripts/map_ec2_to_arc.py new file mode 100644 index 0000000000000..b32cd1c784775 --- /dev/null +++ b/.github/scripts/map_ec2_to_arc.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +"""Map EC2 runner labels to ARC equivalents using .github/arc.yaml. + +Takes a GitHub Actions test matrix, replaces each runner with its ARC +equivalent, and prints the updated matrix as JSON. + +Usage: + python map_ec2_to_arc.py --prefix mt- '{ include: [ + { config: "default", shard: 1, num_shards: 5, runner: "mt-linux.4xlarge" }, + ]}' +""" + +import argparse +import json +import os +import sys +from pathlib import Path + +import yaml + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Map EC2 runner labels to ARC runner labels in a test matrix" + ) + parser.add_argument( + "matrix", + help="GitHub Actions test matrix string to transform", + ) + parser.add_argument( + "--prefix", + default="", + help="Runner prefix to strip from labels (e.g. 'mt-')", + ) + return parser.parse_args() + + +def strip_prefix(label: str, prefix: str) -> str: + if prefix and label.startswith(prefix): + return label[len(prefix) :] + return label + + +def load_mapping(arc_yaml: Path) -> dict[str, str]: + with open(arc_yaml) as f: + data = yaml.safe_load(f) + return data["runner_mapping"] + + +def set_output(name: str, val: str) -> None: + print(f"Setting {name}={val}") + github_output = os.getenv("GITHUB_OUTPUT") + if github_output: + with open(github_output, "a") as f: + print(f"{name}={val}", file=f) + + +def main() -> None: + args = parse_args() + arc_yaml = Path(__file__).resolve().parent.parent / "arc.yaml" + mapping = load_mapping(arc_yaml) + + matrix = yaml.safe_load(args.matrix) + if not matrix: + set_output("test-matrix", args.matrix) + return + + entries = matrix.get("include", []) + if not entries: + set_output("test-matrix", json.dumps(matrix)) + return + + # TODO(huydo): onnxruntime uses hardware_concurrency() to size its thread + # pool, which sees all host CPUs (e.g., 192) on ARC k8s instead of the + # container's cpuset (e.g., 16). This causes pthread_setaffinity_np errors. + # Skip onnx tests on ARC until the onnxruntime session options are fixed to + # use cgroup-aware CPU counts. + excluded_configs = {"onnx"} + filtered = [] + for entry in entries: + if entry.get("config") in excluded_configs: + print(f"Excluding config '{entry['config']}' from ARC test matrix") + continue + filtered.append(entry) + matrix["include"] = filtered + + for entry in filtered: + if "runner" not in entry: + continue + clean = strip_prefix(entry["runner"].strip(), args.prefix) + if clean not in mapping: + print(f"error: no ARC runner found for '{clean}'", file=sys.stderr) + sys.exit(1) + entry["runner"] = args.prefix + mapping[clean] + + set_output("test-matrix", json.dumps(matrix)) + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/runner_determinator.py b/.github/scripts/runner_determinator.py index da57e5f668d3b..40895b84c488f 100644 --- a/.github/scripts/runner_determinator.py +++ b/.github/scripts/runner_determinator.py @@ -1,9 +1,5 @@ # flake8: noqa: G004 -# Note: Copies of this script in runner_determinator.py and _runner-determinator.yml -# must be kept in sync. You can do it easily by running the following command: -# python .github/scripts/update_runner_determinator.py - """ This runner determinator is used to determine which set of runners to run a GitHub job on. It uses the first comment of a GitHub issue (by default @@ -362,11 +358,8 @@ def parse_settings_from_text(settings_text: str) -> Settings: """ try: if settings_text: - # Escape the backtick as well so that we can have the settings in a code block on the GH issue - # for easy reading - # Note: Using ascii for the backtick so that the cat step in _runner-determinator.yml doesn't choke on - # the backtick character in shell commands. - backtick = chr(96) # backtick character + # Strip backticks so settings can be in a code block on the GH issue + backtick = chr(96) settings_text = settings_text.strip(f"\r\n\t{backtick} ") settings = load_yaml(settings_text) diff --git a/.github/scripts/test_filter_test_configs.py b/.github/scripts/test_filter_test_configs.py index 26e38828b7865..1463f82f61eda 100755 --- a/.github/scripts/test_filter_test_configs.py +++ b/.github/scripts/test_filter_test_configs.py @@ -376,13 +376,31 @@ def test_set_periodic_modes(self) -> None: scheduled_test_matrix = set_periodic_modes(test_matrix, job_name) expected_modes = [ - m for m, c in SUPPORTED_PERIODICAL_MODES.items() if c(job_name) + m for m, c in SUPPORTED_PERIODICAL_MODES.items() if c(job_name, None) ] self.assertEqual( len(test_matrix["include"]) * len(expected_modes), len(scheduled_test_matrix["include"]), ) + def test_set_periodic_modes_gpu_runner(self) -> None: + """Job name without 'cuda' but runner indicates a CUDA job (e.g. inductor-unittest).""" + test_matrix = yaml.safe_load( + "{include: [" + '{config: "inductor", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu"}, ' + '{config: "inductor", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu"}' + "]}" + ) + scheduled = set_periodic_modes(test_matrix, "inductor-build / build") + + modes_per_config = [ + entry.get("mem_leak_check") or entry.get("rerun_disabled_tests") + for entry in scheduled["include"] + ] + self.assertIn("mem_leak_check", modes_per_config) + self.assertIn("rerun_disabled_tests", modes_per_config) + self.assertEqual(len(scheduled["include"]), 4) + @mock.patch("filter_test_configs.download_json") def test_remove_disabled_jobs(self, mock_download_json: Any) -> None: mock_download_json.return_value = MOCKED_DISABLED_UNSTABLE_JOBS diff --git a/.github/scripts/test_map_ec2_to_arc.py b/.github/scripts/test_map_ec2_to_arc.py new file mode 100644 index 0000000000000..4414532be546e --- /dev/null +++ b/.github/scripts/test_map_ec2_to_arc.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 + +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path + + +SCRIPT = Path(__file__).resolve().parent / "map_ec2_to_arc.py" + + +def run( + matrix: str, prefix: str = "", github_output: str | None = None +) -> subprocess.CompletedProcess: + cmd = [sys.executable, str(SCRIPT)] + if prefix: + cmd += ["--prefix", prefix] + cmd.append(matrix) + + env = os.environ.copy() + if github_output is not None: + env["GITHUB_OUTPUT"] = github_output + else: + env.pop("GITHUB_OUTPUT", None) + + return subprocess.run(cmd, capture_output=True, text=True, env=env) + + +def parse_output(stdout: str) -> dict: + """Extract the JSON matrix from the 'Setting test-matrix=...' line.""" + prefix = "Setting test-matrix=" + for line in stdout.splitlines(): + if line.startswith(prefix): + return json.loads(line[len(prefix) :]) + raise ValueError(f"no test-matrix output found in: {stdout}") + + +def check(condition: bool, msg: str = "") -> None: + if not condition: + raise AssertionError(msg) + + +def test_basic_matrix(): + matrix = """{ include: [ + { config: "default", shard: 1, num_shards: 1, runner: "linux.4xlarge" }, + { config: "openreg", shard: 1, num_shards: 1, runner: "linux.2xlarge" }, + ]}""" + result = run(matrix) + check(result.returncode == 0, result.stderr) + output = parse_output(result.stdout) + runners = [e["runner"] for e in output["include"]] + check(runners == ["l-x86iavx512-16-128", "l-x86iavx512-8-64"]) + + +def test_matrix_with_prefix(): + matrix = """{ include: [ + { config: "default", shard: 1, num_shards: 7, runner: "mt-linux.4xlarge" }, + { config: "default", shard: 2, num_shards: 7, runner: "mt-linux.4xlarge" }, + { config: "openreg", shard: 1, num_shards: 1, runner: "mt-linux.2xlarge" }, + ]}""" + result = run(matrix, prefix="mt-") + check(result.returncode == 0, result.stderr) + output = parse_output(result.stdout) + runners = [e["runner"] for e in output["include"]] + check( + runners + == [ + "mt-l-x86iavx512-16-128", + "mt-l-x86iavx512-16-128", + "mt-l-x86iavx512-8-64", + ] + ) + + +def test_matrix_without_prefix_when_none_present(): + matrix = """{ include: [ + { config: "default", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" }, + ]}""" + result = run(matrix) + check(result.returncode == 0, result.stderr) + output = parse_output(result.stdout) + check(output["include"][0]["runner"] == "l-x86aavx2-29-113-a10g") + + +def test_unknown_runner_fails(): + matrix = """{ include: [ + { config: "default", shard: 1, num_shards: 1, runner: "bogus.runner" }, + ]}""" + result = run(matrix) + check(result.returncode == 1) + check("no ARC runner found for 'bogus.runner'" in result.stderr) + + +def test_prefix_not_present_on_runner(): + """When --prefix is given but a runner doesn't have it, the raw label is used.""" + matrix = """{ include: [ + { config: "default", shard: 1, num_shards: 1, runner: "linux.4xlarge" }, + ]}""" + result = run(matrix, prefix="mt-") + check(result.returncode == 0, result.stderr) + output = parse_output(result.stdout) + check(output["include"][0]["runner"] == "mt-l-x86iavx512-16-128") + + +def test_preserves_non_runner_fields(): + matrix = """{ include: [ + { config: "default", shard: 3, num_shards: 7, runner: "linux.large" }, + ]}""" + result = run(matrix) + check(result.returncode == 0, result.stderr) + entry = parse_output(result.stdout)["include"][0] + check(entry["config"] == "default") + check(entry["shard"] == 3) + check(entry["num_shards"] == 7) + check(entry["runner"] == "l-x86iavx512-2-4") + + +def test_empty_include_passes_through(): + matrix = """{ include: [] }""" + result = run(matrix) + check(result.returncode == 0, result.stderr) + output = parse_output(result.stdout) + check(output == {"include": []}, f"expected empty include, got {output}") + + +def test_empty_string_passes_through(): + result = run("") + check(result.returncode == 0, result.stderr) + + +def test_mixed_runners(): + matrix = """{ include: [ + { config: "default", shard: 1, num_shards: 1, runner: "linux.4xlarge" }, + { config: "gpu", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" }, + { config: "arm", shard: 1, num_shards: 1, runner: "linux.arm64.2xlarge" }, + ]}""" + result = run(matrix) + check(result.returncode == 0, result.stderr) + output = parse_output(result.stdout) + runners = [e["runner"] for e in output["include"]] + check( + runners + == [ + "l-x86iavx512-16-128", + "l-x86aavx2-29-113-a10g", + "l-arm64g2-6-32", + ] + ) + + +def test_github_output_file(): + """When GITHUB_OUTPUT is set, the script writes test-matrix to that file.""" + matrix = """{ include: [ + { config: "default", shard: 1, num_shards: 1, runner: "linux.2xlarge" }, + ]}""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + tmp_path = f.name + + try: + result = run(matrix, github_output=tmp_path) + check(result.returncode == 0, result.stderr) + + contents = Path(tmp_path).read_text() + check( + contents.startswith("test-matrix="), f"unexpected file contents: {contents}" + ) + written = json.loads(contents[len("test-matrix=") :].strip()) + check(written["include"][0]["runner"] == "l-x86iavx512-8-64") + finally: + os.unlink(tmp_path) + + +if __name__ == "__main__": + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")] + failed = 0 + for t in tests: + try: + t() + print(f" PASS {t.__name__}") + except AssertionError as e: + print(f" FAIL {t.__name__}: {e}") + failed += 1 + if failed: + print(f"\n{failed}/{len(tests)} tests failed") + sys.exit(1) + print(f"\nAll {len(tests)} tests passed") diff --git a/.github/scripts/trymerge.py b/.github/scripts/trymerge.py index 761af8b691d9b..658728b096d47 100755 --- a/.github/scripts/trymerge.py +++ b/.github/scripts/trymerge.py @@ -2034,9 +2034,13 @@ def validate_revert( # For some reason, one can not be a member of private repo, only CONTRIBUTOR if pr.is_base_repo_private(): allowed_reverters.append("CONTRIBUTOR") - # Special case the pytorch-auto-revert app, whose does not have association - # But should be able to issue revert command - if comment.author_url == "https://github.com/apps/pytorch-auto-revert": + # Special case GitHub Apps that don't have a repo association + # but should be able to issue revert commands + allowed_apps = { + "https://github.com/apps/pytorch-auto-revert", + "https://github.com/apps/facebook-github-tools", + } + if comment.author_url in allowed_apps: allowed_reverters.append("NONE") if author_association not in allowed_reverters: @@ -2174,6 +2178,13 @@ def try_revert( f"Failed to fetch dependent PRs: {str(e)}, fall over to single revert" ) + if not shas_and_prs: + raise RuntimeError( + f"No revertable PRs found in ghstack for #{pr.pr_num}. " + f"This typically means the PR is still open (not merged) or " + f"its GitHub state is inconsistent. Only closed/merged PRs can be reverted." + ) + do_revert_prs( repo, pr, diff --git a/.github/scripts/update_runner_determinator.py b/.github/scripts/update_runner_determinator.py deleted file mode 100755 index 772df87c6405a..0000000000000 --- a/.github/scripts/update_runner_determinator.py +++ /dev/null @@ -1,31 +0,0 @@ -#!/usr/bin/env python3 - -import re - - -# Read the contents of runner_determinator.py -with open(".github/scripts/runner_determinator.py") as script_file: - script_content = script_file.read() - -# Indent the script content by 10 spaces to match destination indentation -indented_script_content = "\n".join( - [" " * 10 + line if line else line for line in script_content.splitlines()] -) - -# Read the contents of _runner-determinator.yml -with open(".github/workflows/_runner-determinator.yml") as yml_file: - yml_content = yml_file.read() - -# Replace the content between the markers -new_yml_content = re.sub( - r"(cat < runner_determinator.py\n)(.*?)(\n\s+EOF)", - lambda match: match.group(1) + indented_script_content + match.group(3), - yml_content, - flags=re.DOTALL, -) - -# Save the modified content back to _runner-determinator.yml -with open(".github/workflows/_runner-determinator.yml", "w") as yml_file: - yml_file.write(new_yml_content) - -print("Updated _runner-determinator.yml with the contents of runner_determinator.py") diff --git a/.github/scripts/windows/build_magma.bat b/.github/scripts/windows/build_magma.bat index 75c916ecdbef7..ddb38c9ed9e6d 100644 --- a/.github/scripts/windows/build_magma.bat +++ b/.github/scripts/windows/build_magma.bat @@ -35,6 +35,9 @@ cd magma mkdir build && cd build set GPU_TARGET=All +if "%CUVER_NODOT%" == "132" ( + set CUDA_ARCH_LIST=-gencode=arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_100,code=sm_100 -gencode arch=compute_120,code=sm_120 +) if "%CUVER_NODOT%" == "130" ( set CUDA_ARCH_LIST=-gencode=arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90 -gencode arch=compute_100,code=sm_100 -gencode arch=compute_120,code=sm_120 ) diff --git a/.github/templates/linux_binary_build_workflow.yml.j2 b/.github/templates/linux_binary_build_workflow.yml.j2 index 21a978f801116..4182ae6c33874 100644 --- a/.github/templates/linux_binary_build_workflow.yml.j2 +++ b/.github/templates/linux_binary_build_workflow.yml.j2 @@ -79,7 +79,7 @@ jobs: timeout-minutes: 420 {%- elif config["gpu_arch_type"] == "rocm" %} runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 {%- elif "conda" in build_environment and config["gpu_arch_type"] == "cuda" %} runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.24xlarge.ephemeral @@ -183,7 +183,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 diff --git a/.github/workflows/_bazel-build-test.yml b/.github/workflows/_bazel-build-test.yml index eaebce92ba898..7ebd0b89f5f34 100644 --- a/.github/workflows/_bazel-build-test.yml +++ b/.github/workflows/_bazel-build-test.yml @@ -32,13 +32,33 @@ on: type: string default: "linux.large" description: Runner type + use-arc: + required: false + type: boolean + default: false + description: If true, use ARC (OSDC) runner path instead of EC2. + runner_prefix: + required: false + default: "" + type: string + description: Prefix for runner label + python-version: + required: false + type: string + default: "" + description: Python version to use for the OSDC build. + compiler: + required: false + type: string + default: "" + description: Compiler to use for the OSDC build. env: GIT_DEFAULT_BRANCH: ${{ github.event.repository.default_branch }} jobs: filter: - if: github.repository_owner == 'pytorch' + if: github.repository_owner == 'pytorch' && !inputs.use-arc runs-on: ${{ inputs.runner }} outputs: test-matrix: ${{ steps.filter.outputs.test-matrix }} @@ -62,7 +82,7 @@ jobs: build-and-test: needs: filter # Don't run on forked repos. - if: github.repository_owner == 'pytorch' && needs.filter.outputs.is-test-matrix-empty == 'False' + if: github.repository_owner == 'pytorch' && needs.filter.outputs.is-test-matrix-empty == 'False' && !inputs.use-arc strategy: matrix: ${{ fromJSON(needs.filter.outputs.test-matrix) }} fail-fast: false @@ -73,12 +93,8 @@ jobs: with: github-secret: ${{ secrets.GITHUB_TOKEN }} - # [see note: pytorch repo ref] - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - - name: Setup Linux - uses: ./.github/actions/setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main - name: Login to ECR uses: ./.github/actions/ecr-login @@ -140,7 +156,7 @@ jobs: SHA1: ${{ github.event.pull_request.head.sha || github.sha }} SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2 SCCACHE_REGION: us-east-1 - TORCH_CUDA_ARCH_LIST: 5.2 + TORCH_CUDA_ARCH_LIST: 7.5 DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }} OUR_GITHUB_JOB_ID: ${{ steps.get-job-id.outputs.job-id }} CUDA_VERSION: ${{ inputs.cuda-version }} @@ -213,3 +229,97 @@ jobs: - name: Teardown Linux uses: pytorch/test-infra/.github/actions/teardown-linux@main if: always() + + build-and-test-osdc: + if: github.repository_owner == 'pytorch' && inputs.use-arc + permissions: + id-token: write + contents: read + actions: read + runs-on: ${{ inputs.runner_prefix }}${{ startsWith(inputs.runner, 'l-') && inputs.runner || 'l-x86iavx512-16-128' }} + container: + image: ghcr.io/pytorch/${{ inputs.docker-image-name }} + steps: + - name: Setup Linux + id: setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main + with: + use-arc: true + python-version: ${{ inputs.python-version }} + compiler: ${{ inputs.compiler }} + cuda-version: ${{ inputs.cuda-version }} + github-token: ${{ secrets.GITHUB_TOKEN }} + + - name: Configure AWS credentials + id: aws-creds + continue-on-error: true + uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 + with: + role-to-assume: arn:aws:iam::308535385114:role/arc + aws-region: us-east-1 + role-duration-seconds: 18000 + + - name: Configure Bazel for OSDC pypi-cache + shell: bash + run: | + # On OSDC, cache-enforcer iptables rules TCP-reset direct connections + # to files.pythonhosted.org and pypi.org. + # + # Two download paths need fixing: + # 1) pip_parse's internal deps use http_archive (Bazel's Java downloader) + # -> fix with --experimental_downloader_config URL rewrite + # 2) User pip packages use pip subprocess + # -> fix with --repo_env=PIP_INDEX_URL (envsubst in rules_python 0.31.0) + # + # Bazel strips the scheme before matching rewrite patterns, so the + # pattern is scheme-less. Starting the replacement with http:// makes + # Bazel use that scheme instead of the original https://. + CACHE="http://pypi-cache-cpu.pypi-cache.svc.cluster.local:8080" + cat > /tmp/osdc_downloader.cfg << DLEOF + rewrite files\.pythonhosted\.org/(.*) ${CACHE}/\$1 + rewrite pypi\.org/(.*) ${CACHE}/\$1 + DLEOF + echo "common --experimental_downloader_config=/tmp/osdc_downloader.cfg" >> .bazelrc + echo "common --repo_env=PIP_INDEX_URL=${PIP_INDEX_URL}" >> .bazelrc + echo "common --repo_env=PIP_TRUSTED_HOST=${PIP_TRUSTED_HOST}" >> .bazelrc + + - name: Build + env: + BUILD_ENVIRONMENT: ${{ inputs.build-environment }} + BRANCH: ${{ steps.setup-linux.outputs.branch }} + PR_NUMBER: ${{ github.event.pull_request.number }} + SHA1: ${{ github.event.pull_request.head.sha || github.sha }} + SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2 + SCCACHE_REGION: us-east-1 + SCCACHE_S3_NO_CREDENTIALS: ${{ steps.aws-creds.outcome != 'success' && 'true' || 'false' }} + OUR_GITHUB_JOB_ID: ${{ steps.setup-linux.outputs.job-id }} + CUDA_VERSION: ${{ inputs.cuda-version }} + SKIP_SCCACHE_INITIALIZATION: 1 + TORCH_CUDA_ARCH_LIST: 7.5 + shell: bash + run: | + .ci/pytorch/build.sh + + - name: Test + id: test + timeout-minutes: 120 + env: + BUILD_ENVIRONMENT: ${{ inputs.build-environment }} + BRANCH: ${{ steps.setup-linux.outputs.branch }} + JOB_ID: ${{ steps.setup-linux.outputs.job-id }} + SHA1: ${{ github.event.pull_request.head.sha || github.sha }} + SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2 + SCCACHE_REGION: us-east-1 + SCCACHE_S3_NO_CREDENTIALS: ${{ steps.aws-creds.outcome != 'success' && 'true' || 'false' }} + CUDA_VERSION: ${{ inputs.cuda-version }} + SKIP_SCCACHE_INITIALIZATION: 1 + SHARD_NUMBER: 0 + shell: bash + run: | + .ci/pytorch/test.sh && cp -Lr ./bazel-testlogs ./test/test-reports + + - name: Upload test artifacts + uses: pytorch/pytorch/.github/actions/upload-test-artifacts@main + if: always() && steps.test.conclusion && steps.test.conclusion != 'skipped' + with: + file-suffix: bazel-${{ github.job }}_${{ steps.setup-linux.outputs.job-id }} diff --git a/.github/workflows/_binary-build-flash-attention-wheel-linux.yml b/.github/workflows/_binary-build-flash-attention-wheel-linux.yml index 3fdc1dc4175c9..aa945af50d39c 100644 --- a/.github/workflows/_binary-build-flash-attention-wheel-linux.yml +++ b/.github/workflows/_binary-build-flash-attention-wheel-linux.yml @@ -88,14 +88,11 @@ jobs: github-secret: ${{ secrets.GITHUB_TOKEN }} fail-silently: false - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main + - name: Setup Linux + uses: pytorch/pytorch/.github/actions/setup-linux@main with: submodules: true - - name: Setup Linux - uses: ./.github/actions/setup-linux - - name: Pull Docker image uses: pytorch/test-infra/.github/actions/pull-docker-image@main with: diff --git a/.github/workflows/_binary-build-flash-attention-wheel-windows.yml b/.github/workflows/_binary-build-flash-attention-wheel-windows.yml index 4fc1dc8a53367..5339753a9526f 100644 --- a/.github/workflows/_binary-build-flash-attention-wheel-windows.yml +++ b/.github/workflows/_binary-build-flash-attention-wheel-windows.yml @@ -90,7 +90,7 @@ jobs: echo "CUDA_HOME=${CUDA_PATH}" >> "${GITHUB_ENV}" echo "${CUDA_PATH}/bin" >> "${GITHUB_PATH}" - name: Setup MSVC - uses: ilammy/msvc-dev-cmd@v1 + uses: ilammy/msvc-dev-cmd@0b201ec74fa43914dc39ae48a89fd1d8cb592756 # v1 - name: Remove link.exe conflict run: rm -f /usr/bin/link diff --git a/.github/workflows/_binary-build-linux.yml b/.github/workflows/_binary-build-linux.yml index 461a68eab7bb6..2db2bb75bf978 100644 --- a/.github/workflows/_binary-build-linux.yml +++ b/.github/workflows/_binary-build-linux.yml @@ -154,7 +154,7 @@ jobs: - name: Setup Linux if: inputs.build_environment != 'linux-s390x-binary-manywheel' - uses: ./.github/actions/setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main - name: Login to ECR if: inputs.build_environment != 'linux-s390x-binary-manywheel' diff --git a/.github/workflows/_binary-test-linux.yml b/.github/workflows/_binary-test-linux.yml index ec2bd6ffe5017..4dc035055955e 100644 --- a/.github/workflows/_binary-test-linux.yml +++ b/.github/workflows/_binary-test-linux.yml @@ -138,7 +138,7 @@ jobs: - name: Setup Linux if: inputs.build_environment != 'linux-s390x-binary-manywheel' - uses: ./.github/actions/setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main - name: Login to ECR if: inputs.build_environment != 'linux-s390x-binary-manywheel' diff --git a/.github/workflows/_docs.yml b/.github/workflows/_docs.yml index f5cb186c6f189..d13a69a34dc58 100644 --- a/.github/workflows/_docs.yml +++ b/.github/workflows/_docs.yml @@ -47,6 +47,26 @@ on: description: prefix for runner label type: string default: "" + use-arc: + required: false + type: boolean + default: false + description: If true, use ARC (OSDC) runner path instead of EC2. + python-version: + required: false + type: string + default: "" + description: Python version to use for the OSDC build. + compiler: + required: false + type: string + default: "" + description: Compiler to use for the OSDC build. + cuda-version: + required: false + type: string + default: "" + description: CUDA version to use for the OSDC build. secrets: GH_PYTORCHBOT_TOKEN: required: false @@ -55,7 +75,7 @@ on: jobs: build-docs: # Don't run on forked repos. - if: github.repository_owner == 'pytorch' + if: github.repository_owner == 'pytorch' && !inputs.use-arc runs-on: ${{ matrix.runner }} environment: ${{ (github.ref == 'refs/heads/main' || startsWith(github.event.ref, 'refs/tags/v')) && 'pytorchbot-env' || '' }} strategy: @@ -89,12 +109,8 @@ jobs: To start Python docs build type: cd docs && make html && make coverage - # [see note: pytorch repo ref] - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - - name: Setup Linux - uses: ./.github/actions/setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main - name: Login to ECR uses: ./.github/actions/ecr-login @@ -204,6 +220,40 @@ jobs: path: cppdocs/ s3-prefix: pytorch/pytorch/${{ github.event.pull_request.number }}/cppdocs + - name: Post C++ Docs Coverage Comment + if: ${{ github.event_name == 'pull_request' && matrix.docs_type == 'cpp' && steps.build-docs.outcome == 'success' }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + PR_NUM="${{ github.event.pull_request.number }}" + MARKER="" + + # Only post if the PR touches docs/cpp + if ! gh pr diff "$PR_NUM" --name-only | grep -q '^docs/cpp/'; then + echo "No changes to docs/cpp/, skipping coverage comment." + exit 0 + fi + + # Skip if we already posted + if gh pr view "$PR_NUM" --json comments --jq '.comments[].body' | grep -q "$MARKER"; then + echo "Coverage comment already posted, skipping." + exit 0 + fi + + body="${MARKER} + ## C++ Docs + + - [Doc Preview](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/index.html) + - [API Coverage Report](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/cpp_coverage.txt) + - [HTML Issues Report](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/cpp_html_issues.txt)" + + if [ -f cppdocs/_coverage/index.html ]; then + body="${body} + - [Coverxygen Report (interactive)](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/index.html)" + fi + + gh pr comment "$PR_NUM" --body "$body" + - name: Upload C++ Docs Preview (nightly dry-run) if: ${{ !inputs.push && github.event_name != 'pull_request' && matrix.docs_type == 'cpp' && steps.build-docs.outcome == 'success' }} run: | @@ -215,3 +265,157 @@ jobs: - name: Teardown Linux uses: pytorch/test-infra/.github/actions/teardown-linux@main if: always() + + build-docs-osdc: + if: github.repository_owner == 'pytorch' && inputs.use-arc + permissions: + id-token: write + contents: read + actions: read + runs-on: ${{ matrix.runner }} + container: + image: ${{ inputs.docker-image }} + strategy: + fail-fast: false + matrix: + include: + - docs_type: cpp + runner: ${{ inputs.runner_prefix }}l-x86iavx512-48-384 + timeout-minutes: 360 + - docs_type: python + runner: ${{ inputs.runner_prefix }}l-x86iavx512-16-128 + timeout-minutes: 30 + name: build-docs-${{ matrix.docs_type }}-${{ inputs.push }} + steps: + - name: Setup Linux + id: setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main + with: + use-arc: true + python-version: ${{ inputs.python-version }} + compiler: ${{ inputs.compiler }} + cuda-version: ${{ inputs.cuda-version }} + submodules: false + github-token: ${{ secrets.GITHUB_TOKEN }} + + - name: Configure AWS credentials + id: aws-creds + continue-on-error: true + uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 + with: + role-to-assume: arn:aws:iam::308535385114:role/arc + aws-region: us-east-1 + role-duration-seconds: 18000 + + - name: Download build artifacts + uses: pytorch/pytorch/.github/actions/download-build-artifacts@main + with: + name: ${{ inputs.build-environment }} + s3-bucket: ${{ inputs.s3-bucket }} + use-gha: ${{ steps.aws-creds.outcome != 'success' }} + + - name: Generate netrc (only for docs-push) + if: inputs.push + env: + GITHUB_PYTORCHBOT_TOKEN: ${{ secrets.GH_PYTORCHBOT_TOKEN }} + run: | + rm -rf "${HOME}/.netrc" + echo "machine github.com" > "${HOME}/.netrc" + echo "login pytorchbot" >> "${HOME}/.netrc" + echo "password ${GITHUB_PYTORCHBOT_TOKEN}" >> "${HOME}/.netrc" + + - name: Build ${{ matrix.docs_type }} docs + timeout-minutes: ${{ matrix.timeout-minutes }} + id: build-docs + env: + WITH_PUSH: ${{ inputs.push }} + DOCS_TYPE: ${{ matrix.docs_type }} + RUN_DOXYGEN: ${{ inputs.run-doxygen }} + BUILD_ENVIRONMENT: ${{ inputs.build-environment }} + SHA1: ${{ github.sha }} + shell: bash + run: | + set -ex + if [[ "${GITHUB_REF}" =~ ^refs/tags/v([0-9]+\.[0-9]+)\.* ]]; then + target="${BASH_REMATCH[1]}" + else + target="main" + fi + export DOCS_VERSION="${target}" + pip install $(echo dist/*.whl)[opt-einsum] + ./.ci/pytorch/${DOCS_TYPE}_doc_push_script.sh + + - name: Upload Python Docs Preview + uses: seemethere/upload-artifact-s3@baba72d0712b404f646cebe0730933554ebce96a # v5.1.0 + if: ${{ steps.aws-creds.outcome == 'success' && github.event_name == 'pull_request' && matrix.docs_type == 'python' && steps.build-docs.outcome == 'success' }} + with: + retention-days: 14 + s3-bucket: doc-previews + if-no-files-found: error + path: pytorch_docs/main/ + s3-prefix: pytorch/pytorch/${{ github.event.pull_request.number }} + + - name: Upload C++ Docs Preview + uses: seemethere/upload-artifact-s3@baba72d0712b404f646cebe0730933554ebce96a # v5.1.0 + if: ${{ steps.aws-creds.outcome == 'success' && github.event_name == 'pull_request' && matrix.docs_type == 'cpp' && steps.build-docs.outcome == 'success' }} + with: + retention-days: 14 + if-no-files-found: error + s3-bucket: doc-previews + path: cppdocs/ + s3-prefix: pytorch/pytorch/${{ github.event.pull_request.number }}/cppdocs + + - name: Post C++ Docs Coverage Comment + if: ${{ steps.aws-creds.outcome == 'success' && github.event_name == 'pull_request' && matrix.docs_type == 'cpp' && steps.build-docs.outcome == 'success' }} + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + PR_NUM="${{ github.event.pull_request.number }}" + MARKER="" + + # Only post if the PR touches docs/cpp + if ! gh pr diff "$PR_NUM" --name-only | grep -q '^docs/cpp/'; then + echo "No changes to docs/cpp/, skipping coverage comment." + exit 0 + fi + + # Skip if we already posted + if gh pr view "$PR_NUM" --json comments --jq '.comments[].body' | grep -q "$MARKER"; then + echo "Coverage comment already posted, skipping." + exit 0 + fi + + body="${MARKER} + ## C++ Docs + + - [Doc Preview](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/index.html) + - [API Coverage Report](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/cpp_coverage.txt) + - [HTML Issues Report](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/cpp_html_issues.txt)" + + if [ -f cppdocs/_coverage/index.html ]; then + body="${body} + - [Coverxygen Report (interactive)](https://docs-preview.pytorch.org/pytorch/pytorch/${PR_NUM}/cppdocs/_coverage/index.html)" + fi + + gh pr comment "$PR_NUM" --body "$body" + + - name: Upload C++ Docs Preview (nightly dry-run) + if: ${{ steps.aws-creds.outcome == 'success' && !inputs.push && github.event_name != 'pull_request' && matrix.docs_type == 'cpp' && steps.build-docs.outcome == 'success' }} + run: | + # Unlike EC2 runners, the OSDC container doesn't have aws CLI pre-installed + pip install awscli==1.29.40 + aws s3 cp cppdocs/ s3://doc-previews/pytorch/pytorch/nightly-${{ github.sha }}/cppdocs --recursive --quiet + echo "C++ docs preview available at:" + echo "https://docs-preview.pytorch.org/pytorch/pytorch/nightly-${{ github.sha }}/cppdocs/index.html" + + memory-viz-tests: + # Tests for torch memory visualizer + # tag shangdiy for any issue with this test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-node@v4 + with: + node-version: '20' + - name: Run MemoryViz JS tests (nonretryable) + run: node test/profiler/test_memory_viz.js diff --git a/.github/workflows/_get-changed-files.yml b/.github/workflows/_get-changed-files.yml index f7138a1ddaacc..f3367916f8783 100644 --- a/.github/workflows/_get-changed-files.yml +++ b/.github/workflows/_get-changed-files.yml @@ -25,7 +25,7 @@ jobs: env: GH_TOKEN: ${{ github.token }} run: | - set -e + set -eo pipefail # Check if we're in a pull request context if [ "${{ github.event_name }}" = "pull_request" ] || [ "${{ github.event_name }}" = "pull_request_target" ]; then echo "Running in PR context" @@ -38,8 +38,15 @@ jobs: echo "all_files input is true, returning all files" echo "changed-files=*" >> "$GITHUB_OUTPUT" else - # Use gh CLI to get changed files in the PR with explicit repo - CHANGED_FILES=$(gh api repos/${{ github.repository }}/pulls/$PR_NUMBER/files --paginate --jq '.[] | select(.status != "removed") | .filename' | tr '\n' ' ' | sed 's/ $//') + # Use gh CLI to get changed files in the PR with explicit repo. + # pipefail ensures that if gh fails (e.g. 403 rate limit), the + # error propagates through the pipe instead of being silently + # captured as the file list and injected into downstream lint jobs. + if ! CHANGED_FILES=$(gh api repos/${{ github.repository }}/pulls/$PR_NUMBER/files --paginate --jq '.[] | select(.status != "removed") | .filename' | tr '\n' ' ' | sed 's/ $//'); then + echo "Failed to get changed files from GitHub API, falling back to all files" + echo "changed-files=*" >> "$GITHUB_OUTPUT" + exit 0 + fi # See https://github.com/pytorch/pytorch/pull/134215#issuecomment-2332128790 PYI_FILES_TO_ADD="" diff --git a/.github/workflows/_link_check.yml b/.github/workflows/_link_check.yml index 014e6106b0730..efa5b433947ef 100644 --- a/.github/workflows/_link_check.yml +++ b/.github/workflows/_link_check.yml @@ -11,15 +11,10 @@ on: jobs: lint-urls: if: ${{ github.event_name != 'pull_request' || !contains(github.event.pull_request.labels.*.name, 'skip-url-lint') }} - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: ./.github/workflows/_lint.yml with: - job-name: lint-urls - timeout: 120 - runner: ${{ inputs.runner }}linux.2xlarge - docker-image: ci-image:pytorch-linux-jammy-linter - fetch-depth: 0 - submodules: false - ref: ${{ inputs.ref }} + runner: mt-l-x86iamx-8-16 + docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d script: | ./scripts/lint_urls.sh $( if [ "${{ github.event_name }}" = "pull_request" ]; then @@ -37,15 +32,10 @@ jobs: lint-xrefs: if: ${{ github.event_name != 'pull_request' || !contains(github.event.pull_request.labels.*.name, 'skip-xref-lint') }} - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: ./.github/workflows/_lint.yml with: - job-name: lint-xrefs - timeout: 60 - runner: ${{ inputs.runner }}linux.2xlarge - docker-image: ci-image:pytorch-linux-jammy-linter - fetch-depth: 0 - submodules: false - ref: ${{ inputs.ref }} + runner: mt-l-x86iamx-8-16 + docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d script: | ./scripts/lint_xrefs.sh $( if [ "${{ github.event_name }}" = "pull_request" ]; then diff --git a/.github/workflows/_link_check_osdc.yml b/.github/workflows/_link_check_osdc.yml deleted file mode 100644 index d170d6359cb82..0000000000000 --- a/.github/workflows/_link_check_osdc.yml +++ /dev/null @@ -1,52 +0,0 @@ -on: - workflow_call: - inputs: - runner: - type: string - required: true - ref: - type: string - required: true - -jobs: - lint-urls: - if: ${{ github.event_name != 'pull_request' || !contains(github.event.pull_request.labels.*.name, 'skip-url-lint') }} - uses: ./.github/workflows/_lint.yml - with: - runner: mt-l-x86iamx-8-16 - docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-67eb930 - script: | - ./scripts/lint_urls.sh $( - if [ "${{ github.event_name }}" = "pull_request" ]; then - echo "${{ github.event.pull_request.base.sha }}" "${{ github.event.pull_request.head.sha }}" - else - echo "${{ github.event.before }}" "${{ github.sha }}" - fi - ) || { - echo - echo "URL lint failed." - echo "If this is a transient outage, you can bypass it by adding the \`skip-url-lint\` label to your PR." - echo "Or add \`@lint-ignore\` somewhere on the same line as the URL you want to skip checking." - exit 1 - } - - lint-xrefs: - if: ${{ github.event_name != 'pull_request' || !contains(github.event.pull_request.labels.*.name, 'skip-xref-lint') }} - uses: ./.github/workflows/_lint.yml - with: - runner: mt-l-x86iamx-8-16 - docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-67eb930 - script: | - ./scripts/lint_xrefs.sh $( - if [ "${{ github.event_name }}" = "pull_request" ]; then - echo "${{ github.event.pull_request.base.sha }}" "${{ github.event.pull_request.head.sha }}" - else - echo "${{ github.event.before }}" "${{ github.sha }}" - fi - ) || { - echo - echo "Xref lint failed." - echo "If this is a transient outage, you can bypass it by adding the \`skip-xref-lint\` label to your PR." - echo "Or add \`@lint-ignore\` somewhere on the same line as the reference you want to skip checking." - exit 1 - } diff --git a/.github/workflows/_linux-build.yml b/.github/workflows/_linux-build.yml index c94dbf397be96..14ef0a775204d 100644 --- a/.github/workflows/_linux-build.yml +++ b/.github/workflows/_linux-build.yml @@ -103,34 +103,49 @@ on: required: false type: string default: "" + use-arc: + required: false + type: boolean + default: false + description: If true, use ARC (OSDC) runner path instead of EC2. + python-version: + required: false + type: string + default: "" + description: Python version to use for the OSDC build. + compiler: + required: false + type: string + default: "" + description: Compiler to use for the OSDC build. + cuda-version: + required: false + type: string + default: "" + description: CUDA version to use for the OSDC build. secrets: HUGGING_FACE_HUB_TOKEN: required: false description: | HF Auth token to avoid rate limits when downloading models or datasets from hub - SCRIBE_GRAPHQL_ACCESS_TOKEN: - required: false - description: | - FB app token to write to scribe endpoint outputs: docker-image: - value: ${{ jobs.build.outputs.docker-image }} + value: ${{ jobs.build.outputs.docker-image || jobs.build-osdc.outputs.docker-image }} description: The docker image containing the built PyTorch. test-matrix: - value: ${{ jobs.build.outputs.test-matrix }} + value: ${{ jobs.build.outputs.test-matrix || jobs.build-osdc.outputs.test-matrix }} description: An optional JSON description of what test configs to run later on. build-environment: - value: ${{ jobs.build.outputs.build-environment }} + value: ${{ jobs.build.outputs.build-environment || jobs.build-osdc.outputs.build-environment }} description: Top-level label for what's being built/tested. jobs: build: - environment: ${{ github.ref == 'refs/heads/main' && 'scribe-protected' || startsWith(github.ref, 'refs/heads/release/') && 'scribe-protected' || contains(github.event.pull_request.labels.*.name, 'ci-scribe') && 'scribe-pr' || '' }} # Don't run on forked repos - if: github.repository_owner == 'pytorch' - runs-on: ${{ inputs.runner_prefix}}${{ inputs.runner }} + if: github.repository_owner == 'pytorch' && !inputs.use-arc + runs-on: ${{ inputs.runner_prefix }}${{ inputs.runner }} timeout-minutes: 480 outputs: docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }} @@ -146,19 +161,14 @@ jobs: Build is done inside the container, to start an interactive session run: docker exec -it $(docker container ps --format '{{.ID}}') bash - # [pytorch repo ref] - # Use a pytorch/pytorch reference instead of a reference to the local - # checkout because when we run this action we don't *have* a local - # checkout. In other cases you should prefer a local checkout. - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - with: - no-sudo: true - checkout-mode: treeless - - name: Setup Linux - uses: ./.github/actions/setup-linux - if: inputs.build-environment != 'linux-s390x-binary-manywheel' + id: setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main + with: + python-version: ${{ inputs.python-version }} + compiler: ${{ inputs.compiler }} + cuda-version: ${{ inputs.cuda-version }} + github-token: ${{ secrets.GITHUB_TOKEN }} - name: Login to ECR if: inputs.build-environment != 'linux-s390x-binary-manywheel' @@ -166,13 +176,6 @@ jobs: with: aws-role-to-assume: ${{ inputs.aws-role-to-assume }} - - name: Get workflow job id - id: get-job-id - uses: ./.github/actions/get-workflow-job-id - if: always() - with: - github-token: ${{ secrets.GITHUB_TOKEN }} - - name: Check if can use old whl build id: use-old-whl uses: ./.github/actions/reuse-old-whl @@ -181,8 +184,8 @@ jobs: build-environment: ${{ inputs.build-environment }} run-id: ${{ github.run_id }} github-token: ${{ secrets.GITHUB_TOKEN }} - job-id: ${{ steps.get-job-id.outputs.job-id }} - job-name: ${{ steps.get-job-id.outputs.job-name }} + job-id: ${{ steps.setup-linux.outputs.job-id }} + job-name: ${{ steps.setup-linux.outputs.job-name }} - name: Calculate docker image id: calculate-docker-image @@ -207,10 +210,6 @@ jobs: with: docker-image: ${{ steps.calculate-docker-image.outputs.docker-image }} - - name: Parse ref - id: parse-ref - run: .github/scripts/parse_ref.py - # Apply the filter logic to the build step too if the test-config label is already there - name: Select all requested test configurations (if the test matrix is available) id: filter @@ -219,7 +218,7 @@ jobs: github-token: ${{ secrets.GITHUB_TOKEN }} test-matrix: ${{ inputs.test-matrix }} selected-test-configs: ${{ inputs.selected-test-configs }} - job-name: ${{ steps.get-job-id.outputs.job-name }} + job-name: ${{ steps.setup-linux.outputs.job-name }} - name: Start monitoring script id: monitor-script @@ -227,8 +226,8 @@ jobs: shell: bash continue-on-error: true env: - JOB_ID: ${{ steps.get-job-id.outputs.job-id }} - JOB_NAME: ${{ steps.get-job-id.outputs.job-name }} + JOB_ID: ${{ steps.setup-linux.outputs.job-id }} + JOB_NAME: ${{ steps.setup-linux.outputs.job-name }} WORKFLOW_NAME: ${{ github.workflow }} WORKFLOW_RUN_ID: ${{github.run_id}} MONITOR_LOG_INTERVAL: ${{ inputs.monitor-log-interval }} @@ -242,21 +241,12 @@ jobs: > "../../usage_logs/usage_log_build_${JOB_ID}.txt" 2>&1 & echo "monitor-script-pid=${!}" >> "${GITHUB_OUTPUT}" - - name: Download pytest cache - uses: ./.github/actions/pytest-cache-download - continue-on-error: true - if: inputs.build-environment != 'linux-s390x-binary-manywheel' && steps.use-old-whl.outputs.reuse != 'true' - with: - cache_dir: .pytest_cache - job_identifier: ${{ github.workflow }}_${{ inputs.build-environment }} - s3_bucket: ${{ inputs.s3-bucket }} - - name: Build if: (steps.filter.outputs.is-test-matrix-empty == 'False' || inputs.test-matrix == '') && steps.use-old-whl.outputs.reuse != 'true' id: build env: BUILD_ENVIRONMENT: ${{ inputs.build-environment }} - BRANCH: ${{ steps.parse-ref.outputs.branch }} + BRANCH: ${{ steps.setup-linux.outputs.branch }} PR_NUMBER: ${{ github.event.pull_request.number }} SHA1: ${{ github.event.pull_request.head.sha || github.sha }} # Do not set SCCACHE_S3_KEY_PREFIX to share the cache between all build jobs @@ -268,9 +258,8 @@ jobs: DOCKER_IMAGE: ${{ steps.calculate-docker-image.outputs.docker-image }} DOCKER_IMAGE_S390X: ${{ inputs.docker-image-name }} XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }} - OUR_GITHUB_JOB_ID: ${{ steps.get-job-id.outputs.job-id }} + OUR_GITHUB_JOB_ID: ${{ steps.setup-linux.outputs.job-id }} HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }} - SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }} BUILD_ADDITIONAL_PACKAGES: ${{ inputs.build-additional-packages }} RUNNER: ${{ inputs.runner }} run: | @@ -342,7 +331,6 @@ jobs: -e PR_LABELS \ -e OUR_GITHUB_JOB_ID \ -e HUGGING_FACE_HUB_TOKEN \ - -e SCRIBE_GRAPHQL_ACCESS_TOKEN \ -e BUILD_ADDITIONAL_PACKAGES \ -e RUNNER \ --memory="${TOTAL_AVAILABLE_MEMORY_IN_GB%.*}g" \ @@ -426,7 +414,7 @@ jobs: if-no-files-found: error path: artifacts.zip - - name: copy logs + - name: Copy logs shell: bash if: ${{ always() && steps.build.outcome != 'skipped' && !inputs.disable-monitor && inputs.build-environment != 'linux-s390x-binary-manywheel'}} continue-on-error: true @@ -437,7 +425,7 @@ jobs: - name: Upload raw usage log to s3 if: ${{ always() && steps.build.outcome != 'skipped' && !inputs.disable-monitor && inputs.build-environment != 'linux-s390x-binary-manywheel'}} - uses: seemethere/upload-artifact-s3@v5 + uses: seemethere/upload-artifact-s3@baba72d0712b404f646cebe0730933554ebce96a # v5 with: s3-prefix: | ${{ github.repository }}/${{ github.run_id }}/${{ github.run_attempt }}/artifact @@ -457,12 +445,12 @@ jobs: continue-on-error: true uses: ./.github/actions/upload-utilization-stats with: - job_id: ${{ steps.get-job-id.outputs.job-id }} - job_name: ${{ steps.get-job-id.outputs.job-name }} + job_id: ${{ steps.setup-linux.outputs.job-id }} + job_name: ${{ steps.setup-linux.outputs.job-name }} workflow_name: ${{ github.workflow }} workflow_run_id: ${{github.run_id}} workflow_attempt: ${{github.run_attempt}} - artifact_prefix: usage_log_build_${{ steps.get-job-id.outputs.job-id }} + artifact_prefix: usage_log_build_${{ steps.setup-linux.outputs.job-id }} - name: Teardown Linux uses: pytorch/test-infra/.github/actions/teardown-linux@main @@ -475,3 +463,136 @@ jobs: # on s390x stop the container for clean worker stop docker stop -a || true docker kill -a || true + + build-osdc: + permissions: + id-token: write + contents: read + actions: read + # Don't run on forked repos + if: github.repository_owner == 'pytorch' && inputs.use-arc + runs-on: ${{ inputs.runner_prefix }}${{ startsWith(inputs.runner, 'l-') && inputs.runner || contains(inputs.runner, 'arm64') && 'l-arm64g4-16-62' || 'l-x86iavx512-8-64' }} + container: + image: ghcr.io/pytorch/${{ inputs.docker-image-name }} + timeout-minutes: 480 + outputs: + docker-image: ghcr.io/pytorch/${{ inputs.docker-image-name }} + test-matrix: ${{ steps.map-runners.outputs.test-matrix }} + build-environment: ${{ inputs.build-environment }} + steps: + - name: Setup Linux + id: setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main + with: + use-arc: true + python-version: ${{ inputs.python-version }} + compiler: ${{ inputs.compiler }} + cuda-version: ${{ inputs.cuda-version }} + github-token: ${{ secrets.GITHUB_TOKEN }} + + - name: Check if can use old whl build + id: use-old-whl + uses: pytorch/pytorch/.github/actions/reuse-old-whl@main + if: ${{ inputs.allow-reuse-old-whl }} + with: + build-environment: ${{ inputs.build-environment }} + run-id: ${{ github.run_id }} + github-token: ${{ secrets.GITHUB_TOKEN }} + job-id: ${{ steps.setup-linux.outputs.job-id }} + job-name: ${{ steps.setup-linux.outputs.job-name }} + + # Apply the filter logic to the build step too if the test-config label is already there + - name: Select all requested test configurations (if the test matrix is available) + id: filter + uses: pytorch/pytorch/.github/actions/filter-test-configs@main + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + test-matrix: ${{ inputs.test-matrix }} + selected-test-configs: ${{ inputs.selected-test-configs }} + job-name: ${{ steps.setup-linux.outputs.job-name }} + + - name: Map EC2 runners to ARC runners + id: map-runners + env: + FILTERED_TEST_MATRIX: ${{ steps.filter.outputs.test-matrix }} + RUNNER_PREFIX: ${{ inputs.runner_prefix }} + shell: bash + run: | + python3 .github/scripts/map_ec2_to_arc.py --prefix "${RUNNER_PREFIX}" "${FILTERED_TEST_MATRIX}" + + - name: Configure AWS credentials + id: aws-creds + continue-on-error: true + uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 + with: + role-to-assume: arn:aws:iam::308535385114:role/arc + aws-region: us-east-1 + # The max duration enforced by the server side + role-duration-seconds: 18000 + + - name: Build + if: (steps.filter.outputs.is-test-matrix-empty == 'False' || inputs.test-matrix == '') && steps.use-old-whl.outputs.reuse != 'true' + id: build + env: + BUILD_ENVIRONMENT: ${{ inputs.build-environment }} + BRANCH: ${{ steps.setup-linux.outputs.branch }} + PR_NUMBER: ${{ github.event.pull_request.number }} + SHA1: ${{ github.event.pull_request.head.sha || github.sha }} + SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2 + SCCACHE_REGION: us-east-1 + SCCACHE_S3_NO_CREDENTIALS: ${{ steps.aws-creds.outcome != 'success' && 'true' || 'false' }} + XLA_CLANG_CACHE_S3_BUCKET_NAME: ossci-compiler-clang-cache-circleci-xla + PR_LABELS: ${{ toJson(github.event.pull_request.labels.*.name) }} + TORCH_CUDA_ARCH_LIST: ${{ inputs.cuda-arch-list }} + XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }} + OUR_GITHUB_JOB_ID: ${{ steps.setup-linux.outputs.job-id }} + HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }} + BUILD_ADDITIONAL_PACKAGES: ${{ inputs.build-additional-packages }} + RUNNER: ${{ inputs.runner }} + SKIP_SCCACHE_INITIALIZATION: 1 + shell: bash + run: | + START_TIME=$(date +%s) + .ci/pytorch/build.sh + END_TIME=$(date +%s) + echo "build_time=$((END_TIME - START_TIME))" >> "$GITHUB_OUTPUT" + + - name: Build external packages + id: build-external-packages + if: inputs.build-external-packages != '' && steps.build.outcome != 'skipped' + uses: pytorch/pytorch/.github/actions/build-external-packages@main + with: + build-targets: ${{ inputs.build-external-packages }} + docker-image: ghcr.io/pytorch/${{ inputs.docker-image-name }} + cuda-arch-list: ${{ inputs.cuda-arch-list }} + output-dir: external + + - name: Move external packages to dist + if: steps.build-external-packages.outputs.output_dir != '' && steps.build-external-packages.outcome != 'skipped' + shell: bash + run: | + src="${{ steps.build-external-packages.outputs.output_dir }}" + if [ -d "$src" ]; then + mkdir -p "dist/$(dirname "$src")" + mv "$src" "dist/$(dirname "$src")/" + fi + + - name: Archive artifacts into zip + if: inputs.build-generates-artifacts && steps.build.outcome != 'skipped' && steps.use-old-whl.outputs.reuse != 'true' + run: | + zip -1 -r artifacts.zip dist/ build/custom_test_artifacts build/lib build/bin .additional_ci_files + + - name: Store build artifacts + if: inputs.build-generates-artifacts && (steps.build.outcome != 'skipped' || steps.use-old-whl.outputs.reuse == 'true') + uses: pytorch/pytorch/.github/actions/upload-build-artifacts@main + with: + name: ${{ inputs.build-environment }} + s3-bucket: ${{ inputs.s3-bucket }} + use-gha: ${{ steps.aws-creds.outcome != 'success' }} + + - name: Upload sccache stats + if: steps.build.outcome != 'skipped' && steps.aws-creds.outcome == 'success' + uses: pytorch/pytorch/.github/actions/upload-sccache-stats@main + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + build-time: ${{ steps.build.outputs.build_time }} diff --git a/.github/workflows/_linux-test-stable-fa3.yml b/.github/workflows/_linux-test-stable-fa3.yml index f2e16712ff447..98d64b9246869 100644 --- a/.github/workflows/_linux-test-stable-fa3.yml +++ b/.github/workflows/_linux-test-stable-fa3.yml @@ -54,15 +54,13 @@ jobs: # Don't run on forked repos if: github.repository_owner == 'pytorch' runs-on: linux.aws.h100 - timeout-minutes: ${{ inputs.timeout-minutes || 30 }} + timeout-minutes: ${{ inputs.timeout-minutes || 60 }} permissions: id-token: write contents: read steps: - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - with: - no-sudo: true + - name: Setup Linux + uses: pytorch/pytorch/.github/actions/setup-linux@main - name: Checkout flash-attention as a secondary repository uses: actions/checkout@v4 @@ -70,9 +68,6 @@ jobs: repository: Dao-AILab/flash-attention path: flash-attention - - name: Setup Linux - uses: ./.github/actions/setup-linux - - name: Login to ECR uses: ./.github/actions/ecr-login diff --git a/.github/workflows/_linux-test.yml b/.github/workflows/_linux-test.yml index ae8fadd16e577..65b73e139971d 100644 --- a/.github/workflows/_linux-test.yml +++ b/.github/workflows/_linux-test.yml @@ -72,6 +72,38 @@ on: required: false type: number default: 1 + use-arc: + required: false + type: boolean + default: false + description: If true, use ARC (OSDC) runner path instead of EC2. + python-version: + required: false + type: string + default: "" + description: Python version to use for the OSDC test. + compiler: + required: false + type: string + default: "" + description: Compiler to use for the OSDC test. + cuda-version: + required: false + type: string + default: "" + description: CUDA version to use for the OSDC test. + export-profiler-trace: + description: | + If set to "1", export Chrome profiler traces from performance benchmarks. + required: false + type: string + default: "" + enable-torch-trace: + description: | + If set to "1", enable TORCH_TRACE structured logging and collect tlparse output. + required: false + type: string + default: "" secrets: HUGGING_FACE_HUB_TOKEN: required: false @@ -92,7 +124,7 @@ env: jobs: test: # Don't run on forked repos or empty test matrix - if: github.repository_owner == 'pytorch' && toJSON(fromJSON(inputs.test-matrix).include) != '[]' + if: github.repository_owner == 'pytorch' && toJSON(fromJSON(inputs.test-matrix).include) != '[]' && !inputs.use-arc strategy: matrix: ${{ fromJSON(inputs.test-matrix) }} fail-fast: false @@ -112,23 +144,15 @@ jobs: All testing is done inside the container, to start an interactive session run: docker exec -it $(docker container ps --format '{{.ID}}') bash - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - with: - no-sudo: true - checkout-mode: treeless - submodules: false - - - name: Setup Python - if: contains(matrix.runner, 'b200') - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: '3.12' - cache: pip - - name: Setup Linux - uses: ./.github/actions/setup-linux - if: inputs.build-environment != 'linux-s390x-binary-manywheel' && !contains(matrix.runner, 'b200') + id: setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main + with: + python-version: ${{ inputs.python-version }} + compiler: ${{ inputs.compiler }} + cuda-version: ${{ inputs.cuda-version }} + submodules: 'false' + github-token: ${{ secrets.GITHUB_TOKEN }} - name: Check TPU Availability id: check-tpu @@ -192,21 +216,14 @@ jobs: run: echo "SCCACHE_SERVER_PORT_DOCKER_FLAG=-e SCCACHE_SERVER_PORT=$((RUNNER_UID + 4226))" >> "${GITHUB_ENV}" if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'true' && !contains(matrix.runner, 'b200') }} - - name: Get workflow job id - id: get-job-id - uses: ./.github/actions/get-workflow-job-id - if: always() - with: - github-token: ${{ secrets.GITHUB_TOKEN }} - - name: Start monitoring script id: monitor-script if: ${{ !inputs.disable-monitor }} shell: bash continue-on-error: true env: - JOB_ID: ${{ steps.get-job-id.outputs.job-id }} - JOB_NAME: ${{ steps.get-job-id.outputs.job-name }} + JOB_ID: ${{ steps.setup-linux.outputs.job-id }} + JOB_NAME: ${{ steps.setup-linux.outputs.job-name }} WORKFLOW_NAME: ${{ github.workflow }} WORKFLOW_RUN_ID: ${{github.run_id}} MONITOR_LOG_INTERVAL: ${{ inputs.monitor-log-interval }} @@ -266,10 +283,6 @@ jobs: echo "CUDA libraries:" ls -la win-torch-wheel-extracted/lib/x64/ || echo "No CUDA libraries found" - - name: Parse ref - id: parse-ref - run: .github/scripts/parse_ref.py - - name: Check for keep-going label and re-enabled test issues # This uses the filter-test-configs action because it conveniently # checks for labels and re-enabled test issues. It does not actually do @@ -279,7 +292,7 @@ jobs: with: github-token: ${{ secrets.GITHUB_TOKEN }} test-matrix: ${{ inputs.test-matrix }} - job-name: ${{ steps.get-job-id.outputs.job-name }} + job-name: ${{ steps.setup-linux.outputs.job-name }} - name: Set Test step time id: test-timeout @@ -318,9 +331,9 @@ jobs: GITHUB_RUN_ID: ${{ github.run_id }} GITHUB_RUN_NUMBER: ${{ github.run_number }} GITHUB_RUN_ATTEMPT: ${{ github.run_attempt }} - JOB_ID: ${{ steps.get-job-id.outputs.job-id }} - JOB_NAME: ${{ steps.get-job-id.outputs.job-name }} - BRANCH: ${{ steps.parse-ref.outputs.branch }} + JOB_ID: ${{ steps.setup-linux.outputs.job-id }} + JOB_NAME: ${{ steps.setup-linux.outputs.job-name }} + BRANCH: ${{ steps.setup-linux.outputs.branch }} SHA1: ${{ github.event.pull_request.head.sha || github.sha }} BASE_SHA: ${{ github.event.pull_request.base.sha || github.sha }} TEST_CONFIG: ${{ matrix.config }} @@ -347,6 +360,8 @@ jobs: PYTORCH_TEST_RERUN_DISABLED_TESTS: ${{ matrix.rerun_disabled_tests && '1' || '0' }} TESTS_TO_INCLUDE: ${{ inputs.tests-to-include }} DASHBOARD_TAG: ${{ inputs.dashboard-tag }} + EXPORT_PROFILER_TRACE: ${{ inputs.export-profiler-trace }} + ENABLE_TORCH_TRACE: ${{ inputs.enable-torch-trace }} VLLM_TEST_HUGGING_FACE_TOKEN: ${{ secrets.VLLM_TEST_HUGGING_FACE_TOKEN }} HF_CACHE: /mnt/hf_cache # Use offline mode by default, only enable online mode to refresh the models @@ -358,7 +373,7 @@ jobs: HF_DATASETS_OFFLINE: ${{ (github.event_name == 'schedule' || contains(steps.keep-going.outputs.labels, 'ci-refresh-hf-cache')) && '0' || '1' }} HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }} SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }} - ARTIFACTS_FILE_SUFFIX: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.get-job-id.outputs.job-id }} + ARTIFACTS_FILE_SUFFIX: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.setup-linux.outputs.job-id }} TORCH_TPU: ${{ steps.check-tpu.outputs.has_tpu == 'true' && '1' || '' }} TORCH_TPU_TEXT_FILE: /var/lib/jenkins/workspace/.github/ci_commit_pins/torch_tpu.txt run: | @@ -366,8 +381,6 @@ jobs: if [[ $TEST_CONFIG == 'multigpu' ]]; then TEST_COMMAND=.ci/pytorch/multigpu-test.sh - elif [[ $BUILD_ENVIRONMENT == *onnx* ]]; then - TEST_COMMAND=.ci/onnx/test.sh else TEST_COMMAND=.ci/pytorch/test.sh fi @@ -458,6 +471,8 @@ jobs: -e HF_DATASETS_OFFLINE \ -e SCRIBE_GRAPHQL_ACCESS_TOKEN \ -e DASHBOARD_TAG \ + -e EXPORT_PROFILER_TRACE \ + -e ENABLE_TORCH_TRACE \ -e ARTIFACTS_FILE_SUFFIX \ -e TORCH_TPU \ -e TORCH_TPU_TEXT_FILE \ @@ -540,7 +555,7 @@ jobs: uses: ./.github/actions/upload-test-artifacts if: always() && steps.test.conclusion && steps.test.conclusion != 'skipped' with: - file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.get-job-id.outputs.job-id }} + file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.setup-linux.outputs.job-id }} use-gha: ${{ inputs.use-gha }} s3-bucket: ${{ inputs.s3-bucket }} @@ -564,8 +579,8 @@ jobs: continue-on-error: true uses: ./.github/actions/upload-utilization-stats with: - job_id: ${{ steps.get-job-id.outputs.job-id }} - job_name: ${{ steps.get-job-id.outputs.job-name }} + job_id: ${{ steps.setup-linux.outputs.job-id }} + job_name: ${{ steps.setup-linux.outputs.job-name }} workflow_name: ${{ github.workflow }} workflow_run_id: ${{github.run_id}} workflow_attempt: ${{github.run_attempt}} @@ -581,3 +596,226 @@ jobs: # on s390x stop the container for clean worker stop docker stop -a || true docker kill -a || true + + test-osdc: + # Don't run on forked repos or empty test matrix + if: github.repository_owner == 'pytorch' && toJSON(fromJSON(inputs.test-matrix).include) != '[]' && inputs.use-arc + strategy: + matrix: ${{ fromJSON(inputs.test-matrix) }} + fail-fast: false + environment: ${{ github.ref == 'refs/heads/main' && 'scribe-protected' || startsWith(github.ref, 'refs/heads/release/') && 'scribe-protected' || contains(github.event.pull_request.labels.*.name, 'ci-scribe') && 'scribe-pr' || '' }} + runs-on: ${{ matrix.runner }} + container: + image: ${{ inputs.docker-image }} + options: "--gpus all" + timeout-minutes: ${{ matrix.mem_leak_check == 'mem_leak_check' && 600 || inputs.timeout-minutes }} + permissions: + id-token: write + contents: read + actions: read + steps: + - name: Setup Linux + id: setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main + with: + use-arc: true + python-version: ${{ inputs.python-version }} + compiler: ${{ inputs.compiler }} + cuda-version: ${{ inputs.cuda-version }} + submodules: 'false' + github-token: ${{ secrets.GITHUB_TOKEN }} + + - name: Configure AWS credentials + id: aws-creds + continue-on-error: true + uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 + with: + role-to-assume: arn:aws:iam::308535385114:role/arc + aws-region: us-east-1 + # The max duration enforced by the server side + role-duration-seconds: 18000 + + - name: Download build artifacts + uses: pytorch/pytorch/.github/actions/download-build-artifacts@main + with: + name: ${{ inputs.build-environment }} + s3-bucket: ${{ inputs.s3-bucket }} + use-gha: ${{ steps.aws-creds.outcome != 'success' || inputs.use-gha }} + + - name: Download TD artifacts + continue-on-error: true + uses: pytorch/pytorch/.github/actions/download-td-artifacts@main + + - name: Download Windows torch wheel for cross-compilation + if: matrix.win_torch_wheel_artifact != '' + uses: seemethere/download-artifact-s3@1da556a7aa0a088e3153970611f6c432d58e80e6 # v4.2.0 + with: + name: ${{ matrix.win_torch_wheel_artifact }} + path: win-torch-wheel + + - name: Extract Windows wheel and setup CUDA libraries + if: matrix.win_torch_wheel_artifact != '' + shell: bash + run: | + set -x + + # Find the wheel file + WHEEL_FILE=$(find win-torch-wheel -name "*.whl" -type f | head -n 1) + if [ -z "$WHEEL_FILE" ]; then + echo "Error: No wheel file found in win-torch-wheel directory" + exit 1 + fi + echo "Found wheel file: $WHEEL_FILE" + + # Unzip the wheel file + unzip -q "$WHEEL_FILE" -d win-torch-wheel-extracted + echo "Extracted wheel contents" + + # Setup CUDA libraries (cuda.lib and cudart.lib) directory + mkdir -p win-torch-wheel-extracted/lib/x64 + if [ -f "win-torch-wheel/cuda.lib" ]; then + mv win-torch-wheel/cuda.lib win-torch-wheel-extracted/lib/x64/ + echo "Moved cuda.lib to win-torch-wheel-extracted/lib/x64/" + fi + if [ -f "win-torch-wheel/cudart.lib" ]; then + mv win-torch-wheel/cudart.lib win-torch-wheel-extracted/lib/x64/ + echo "Moved cudart.lib to win-torch-wheel-extracted/lib/x64/" + fi + + # Verify CUDA libraries are present + echo "CUDA libraries:" + ls -la win-torch-wheel-extracted/lib/x64/ || echo "No CUDA libraries found" + + - name: Check for keep-going label and re-enabled test issues + # This uses the filter-test-configs action because it conveniently + # checks for labels and re-enabled test issues. It does not actually do + # any filtering. All filtering is done in the build step. + id: keep-going + uses: pytorch/pytorch/.github/actions/filter-test-configs@main + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + test-matrix: ${{ inputs.test-matrix }} + job-name: ${{ steps.setup-linux.outputs.job-name }} + + - name: Set test step time + id: test-timeout + shell: bash + env: + JOB_TIMEOUT: ${{ matrix.mem_leak_check == 'mem_leak_check' && 600 || inputs.timeout-minutes }} + run: | + echo "timeout=$((JOB_TIMEOUT-30))" >> "${GITHUB_OUTPUT}" + + - name: Test + id: test + timeout-minutes: ${{ fromJson(steps.test-timeout.outputs.timeout) }} + env: + BUILD_ENVIRONMENT: ${{ inputs.build-environment }} + PR_NUMBER: ${{ github.event.pull_request.number }} + GITHUB_REPOSITORY: ${{ github.repository }} + GITHUB_WORKFLOW: ${{ github.workflow }} + GITHUB_JOB: ${{ github.job }} + GITHUB_RUN_ID: ${{ github.run_id }} + GITHUB_RUN_NUMBER: ${{ github.run_number }} + GITHUB_RUN_ATTEMPT: ${{ github.run_attempt }} + JOB_ID: ${{ steps.setup-linux.outputs.job-id }} + JOB_NAME: ${{ steps.setup-linux.outputs.job-name }} + BRANCH: ${{ steps.setup-linux.outputs.branch }} + SHA1: ${{ github.event.pull_request.head.sha || github.sha }} + BASE_SHA: ${{ github.event.pull_request.base.sha || github.sha }} + TEST_CONFIG: ${{ matrix.config }} + SHARD_NUMBER: ${{ matrix.shard }} + NUM_TEST_SHARDS: ${{ matrix.num_shards }} + EXTRA_FLAGS: ${{ matrix.extra_flags || '' }} + OP_BENCHMARK_TESTS: ${{ matrix.op_benchmark_tests }} + REENABLED_ISSUES: ${{ steps.keep-going.outputs.reenabled-issues }} + CONTINUE_THROUGH_ERROR: ${{ steps.keep-going.outputs.keep-going }} + VERBOSE_TEST_LOGS: ${{ steps.keep-going.outputs.ci-verbose-test-logs }} + TEST_SHOWLOCALS: ${{ steps.keep-going.outputs.ci-test-showlocals }} + NO_TEST_TIMEOUT: ${{ steps.keep-going.outputs.ci-no-test-timeout }} + NO_TD: ${{ steps.keep-going.outputs.ci-no-td }} + TD_DISTRIBUTED: ${{ steps.keep-going.outputs.ci-td-distributed }} + SCCACHE_BUCKET: ossci-compiler-cache-circleci-v2 + SCCACHE_REGION: us-east-1 + SCCACHE_S3_NO_CREDENTIALS: ${{ steps.aws-creds.outcome != 'success' && 'true' || 'false' }} + XLA_CUDA: ${{ contains(inputs.build-environment, 'xla') && '0' || '' }} + XLA_CLANG_CACHE_S3_BUCKET_NAME: ossci-compiler-clang-cache-circleci-xla + PYTORCH_TEST_CUDA_MEM_LEAK_CHECK: ${{ matrix.mem_leak_check && '1' || '0' }} + PYTORCH_TEST_RERUN_DISABLED_TESTS: ${{ matrix.rerun_disabled_tests && '1' || '0' }} + TESTS_TO_INCLUDE: ${{ inputs.tests-to-include }} + DASHBOARD_TAG: ${{ inputs.dashboard-tag }} + VLLM_TEST_HUGGING_FACE_TOKEN: ${{ secrets.VLLM_TEST_HUGGING_FACE_TOKEN }} + HF_CACHE: /mnt/hf_cache + TRANSFORMERS_OFFLINE: ${{ (github.event_name == 'schedule' || contains(steps.keep-going.outputs.labels, 'ci-refresh-hf-cache')) && '0' || '1' }} + HF_DATASETS_OFFLINE: ${{ (github.event_name == 'schedule' || contains(steps.keep-going.outputs.labels, 'ci-refresh-hf-cache')) && '0' || '1' }} + HUGGING_FACE_HUB_TOKEN: ${{ secrets.HUGGING_FACE_HUB_TOKEN }} + SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.SCRIBE_GRAPHQL_ACCESS_TOKEN }} + ARTIFACTS_FILE_SUFFIX: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.setup-linux.outputs.job-id }} + TORCH_TPU_TEXT_FILE: /var/lib/jenkins/workspace/.github/ci_commit_pins/torch_tpu.txt + USE_ARC: "1" + shell: bash + run: | + set -x + + if [[ "${TEST_CONFIG}" == 'multigpu' ]]; then + TEST_COMMAND=.ci/pytorch/multigpu-test.sh + else + TEST_COMMAND=.ci/pytorch/test.sh + fi + + # Just create an empty HF_CACHE dir if it doesn't exist. This dir is not + # used for anything besides vLLM jobs + if [[ ! -d "${HF_CACHE}" ]]; then + export HF_CACHE="${RUNNER_TEMP}/hf_cache" + mkdir -p "${HF_CACHE}" + + # When there is no cache directory, e.g. benchmark, the job has no + # way but to reach out to HF if needed + export TRANSFORMERS_OFFLINE=0 + export HF_DATASETS_OFFLINE=0 + fi + + # shellcheck disable=SC2046 + python3 -m pip install $(echo dist/*.whl)[opt-einsum] + ${TEST_COMMAND} + + - name: Configure AWS credentials + id: aws-creds-benchmark + continue-on-error: true + uses: aws-actions/configure-aws-credentials@8df5847569e6427dd6c4fb1cf565c83acfa8afa7 + with: + role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_upload-benchmark-results + aws-region: us-east-1 + + - name: Upload pytest cache if tests failed + uses: pytorch/pytorch/.github/actions/pytest-cache-upload@main + continue-on-error: true + if: failure() && steps.test.conclusion && steps.test.conclusion == 'failure' && steps.aws-creds-benchmark.outcome == 'success' + with: + cache_dir: .pytest_cache + shard: ${{ matrix.shard }} + sha: ${{ github.event.pull_request.head.sha || github.sha }} + test_config: ${{ matrix.config }} + job_identifier: ${{ github.workflow }}_${{ inputs.build-environment }} + + - name: Upload the benchmark results + if: steps.aws-creds-benchmark.outcome == 'success' + uses: pytorch/test-infra/.github/actions/upload-benchmark-results@main + with: + benchmark-results-dir: test/test-reports + dry-run: false + schema-version: v3 + github-token: ${{ secrets.GITHUB_TOKEN }} + + - name: Print remaining test logs + shell: bash + if: always() && steps.test.conclusion + run: | + cat test/**/*_toprint.log || true + + - name: Upload test artifacts + uses: pytorch/pytorch/.github/actions/upload-test-artifacts@main + if: always() && steps.test.conclusion && steps.test.conclusion != 'skipped' + with: + file-suffix: ${{ github.job }}-${{ matrix.config }}-${{ matrix.shard }}-${{ matrix.num_shards }}-${{ matrix.runner }}_${{ steps.setup-linux.outputs.job-id }} + use-gha: ${{ inputs.use-gha }} + s3-bucket: ${{ inputs.s3-bucket }} diff --git a/.github/workflows/_runner-determinator.yml b/.github/workflows/_runner-determinator.yml index ec3d05c316c83..06aefad9dd634 100644 --- a/.github/workflows/_runner-determinator.yml +++ b/.github/workflows/_runner-determinator.yml @@ -7,7 +7,8 @@ on: required: false type: string description: | - List of experiments for this workflow. If not defined, all default experiments are included. + Comma-separated list of non-default experiments to opt into for this workflow. + These are added on top of all default experiments. If not defined, only default experiments are included. opt_out_experiments: required: false type: string @@ -62,685 +63,8 @@ jobs: OPT_OUT_EXPERIMENTS: ${{ inputs.opt_out_experiments }} PR_NUMBER: ${{ github.event.pull_request.number }} steps: - # TODO: Remove the hardcoded step below - # Hardcoding below is temporary for testing ALI runners - # This file below should match the script found in .github/scripts/runner_determinator.py - - name: Hardcode runner-determinator script - id: hardcode-script - run: | - cat < runner_determinator.py - # flake8: noqa: G004 - - # Note: Copies of this script in runner_determinator.py and _runner-determinator.yml - # must be kept in sync. You can do it easily by running the following command: - # python .github/scripts/update_runner_determinator.py - - """ - This runner determinator is used to determine which set of runners to run a - GitHub job on. It uses the first comment of a GitHub issue (by default - https://github.com/pytorch/test-infra/issues/5132) to define the configuration - of which runners should be used to run which job. - - The configuration has two parts, the settings and a list of opted-in users, - separated by a line containing "---". If the line is not present, the - settings are considered to be empty with only the second part, the user - list, defined. - - The first part is a YAML block that defines the rollout settings. This can be - used to define any settings that are needed to determine which runners to use. - It's fields are defined by the RolloutSettings class below. - - The second part is a list of users who are explicitly opted in to the LF fleet. - The user list is also a comma separated list of additional features or - experiments which the user could be opted in to. - - The user list has the following rules: - - - Users are GitHub usernames, which must start with the @ prefix - - Each user is also a comma-separated list of features/experiments to enable - - A "#" prefix opts the user out of all experiments - - Example config: - # A list of experiments that can be opted into. - # This defines the behavior they'll induce when opted into. - # Expected syntax is: - # [experiment_name]: # Name of the experiment. Also used for the label prefix. - # rollout_perc: [int] # % of workflows to run with this experiment when users are not opted in. - - experiments: - lf: - rollout_percent: 25 - all_branches: false - default: true - --- - - # Opt-ins: - # Users can opt into the LF fleet by adding their GitHub username to this list - # and specifying experiments to enable in a comma-separated list. - # To always opt out of an experiment, prefix it with a "-". - # Experiments should be from the above list. - - @User1,-lf,split_build - @User2,lf - @User3,split_build - """ - - import json - import logging - import os - import random - import re - import sys - from argparse import ArgumentParser - from collections.abc import Iterable - from functools import cache - from logging import LogRecord - from typing import Any, NamedTuple - from urllib.request import Request, urlopen - - import yaml - from github import Auth, Github - from github.Issue import Issue - - - DEFAULT_LABEL_PREFIX = "" # use meta runners - WORKFLOW_LABEL_LF = "lf." # use runners from the linux foundation - WORKFLOW_LABEL_LF_CANARY = "lf.c." # use canary runners from the linux foundation - - GITHUB_OUTPUT = os.getenv("GITHUB_OUTPUT", "") - GH_OUTPUT_KEY_AMI = "runner-ami" - GH_OUTPUT_KEY_LABEL_TYPE = "label-type" - GH_OUTPUT_KEY_USE_ARC = "use-arc" - OPT_OUT_LABEL = "no-runner-experiments" - - SETTING_EXPERIMENTS = "experiments" - - LF_FLEET_EXPERIMENT = "lf" - ARC_FLEET_EXPERIMENT = "arc" - CANARY_FLEET_SUFFIX = ".c" - - ARC_LABEL_PREFIX = "mt-" - ARC_CANARY_LABEL_PREFIX = "c-" - - - class Experiment(NamedTuple): - rollout_perc: float = ( - 0 # Percentage of workflows to experiment on when user is not opted-in. - ) - all_branches: bool = ( - False # If True, the experiment is also enabled on the exception branches - ) - default: bool = ( - True # If True, the experiment is enabled by default for all queries - ) - - # Add more fields as needed - - - class RunnerPrefixResult(NamedTuple): - prefix: str - use_arc: bool = False - - - class Settings(NamedTuple): - """ - Settings for the experiments that can be opted into. - """ - - experiments: dict[str, Experiment] = {} - - - class ColorFormatter(logging.Formatter): - """Color codes the log messages based on the log level""" - - COLORS = { - "WARNING": "\033[33m", # Yellow - "ERROR": "\033[31m", # Red - "CRITICAL": "\033[31m", # Red - "INFO": "\033[0m", # Reset - "DEBUG": "\033[0m", # Reset - } - - def format(self, record: LogRecord) -> str: - log_color = self.COLORS.get(record.levelname, "\033[0m") # Default to reset - record.msg = f"{log_color}{record.msg}\033[0m" - return super().format(record) - - - handler = logging.StreamHandler() - handler.setFormatter(ColorFormatter(fmt="%(levelname)-8s: %(message)s")) - - log = logging.getLogger(os.path.basename(__file__)) - log.addHandler(handler) - log.setLevel(logging.INFO) - - - def set_github_output(key: str, value: str) -> None: - """ - Defines outputs of the github action that invokes this script - """ - if not GITHUB_OUTPUT: - # See https://github.blog/changelog/2022-10-11-github-actions-deprecating-save-state-and-set-output-commands/ for deprecation notice - log.warning( - "No env var found for GITHUB_OUTPUT, you must be running this code locally. Falling back to the deprecated print method." - ) - print(f"::set-output name={key}::{value}") - return - - with open(GITHUB_OUTPUT, "a") as f: - log.info(f"Setting output: {key}='{value}'") - f.write(f"{key}={value}\n") - - - def _str_comma_separated_to_set(value: str) -> frozenset[str]: - return frozenset( - filter(lambda itm: itm != "", map(str.strip, value.strip(" \n\t").split(","))) - ) - - - def parse_args() -> Any: - parser = ArgumentParser("Get dynamic rollout settings") - parser.add_argument("--github-token", type=str, required=True, help="GitHub token") - parser.add_argument( - "--github-issue-repo", - type=str, - required=False, - default="pytorch/test-infra", - help="GitHub repo to get the issue", - ) - parser.add_argument( - "--github-repo", - type=str, - required=True, - help="GitHub repo where CI is running", - ) - parser.add_argument( - "--github-issue", type=int, required=True, help="GitHub issue number" - ) - parser.add_argument( - "--github-actor", type=str, required=True, help="GitHub triggering_actor" - ) - parser.add_argument( - "--github-issue-owner", type=str, required=True, help="GitHub issue owner" - ) - parser.add_argument( - "--github-branch", type=str, required=True, help="Current GitHub branch or tag" - ) - parser.add_argument( - "--github-ref-type", - type=str, - required=True, - help="Current GitHub ref type, branch or tag", - ) - parser.add_argument( - "--eligible-experiments", - type=_str_comma_separated_to_set, - required=False, - default="", - help="comma separated list of experiments to check, if omitted all experiments marked with default=True are checked", - ) - parser.add_argument( - "--opt-out-experiments", - type=_str_comma_separated_to_set, - required=False, - default="", - help=( - "comma separated list of experiments to opt-out of. If unset, no opt-outs will occur. " - "If the same experiment is listed both here and in '--eligible-experiments' opt-out will take priority." - ), - ) - parser.add_argument( - "--pr-number", - type=str, - required=False, - default="", - help="the optional PR number where this is run", - ) - - return parser.parse_args() - - - def get_gh_client(github_token: str) -> Github: # type: ignore[no-any-unimported] - auth = Auth.Token(github_token) - return Github(auth=auth) - - - def get_issue(gh: Github, repo: str, issue_num: int) -> Issue: # type: ignore[no-any-unimported] - repo = gh.get_repo(repo) - return repo.get_issue(number=issue_num) - - - def get_potential_pr_author( - github_token: str, repo: str, username: str, ref_type: str, ref_name: str - ) -> str: - # If the trigger was a new tag added by a bot, this is a ciflow case - # Fetch the actual username from the original PR. The PR number is - # embedded in the tag name: ciflow// - - gh = get_gh_client(github_token) - - if username == "pytorch-bot[bot]" and ref_type == "tag": - split_tag = ref_name.split("/") - if ( - len(split_tag) == 3 - and split_tag[0] == "ciflow" - and split_tag[2].isnumeric() - ): - pr_number = split_tag[2] - try: - repository = gh.get_repo(repo) - pull = repository.get_pull(number=int(pr_number)) - except Exception as e: - raise Exception( # noqa: TRY002 - f"issue with pull request {pr_number} from repo {repository}" - ) from e - return pull.user.login # type: ignore[no-any-return] - # In all other cases, return the original input username - return username - - - def is_exception_branch(branch: str) -> bool: - """ - Branches that get opted out of experiments by default, until they're explicitly enabled. - """ - return branch.split("/", maxsplit=1)[0] in { - "main", - "nightly", - "release", - "landchecks", - } - - - def load_yaml(yaml_text: str) -> Any: - try: - data = yaml.safe_load(yaml_text) - return data - except yaml.YAMLError: - log.exception("Error loading YAML") - raise - - - def extract_settings_user_opt_in_from_text(rollout_state: str) -> tuple[str, str]: - """ - Extracts the text with settings, if any, and the opted in users from the rollout state. - - If the issue body contains "---" then the text above that is the settings - and the text below is the list of opted in users. - - If it doesn't contain "---" then the settings are empty and the rest is the users. - """ - rollout_state_parts = rollout_state.split("---") - if len(rollout_state_parts) >= 2: - return rollout_state_parts[0], rollout_state_parts[1] - else: - return "", rollout_state - - - class UserOptins(dict[str, list[str]]): - """ - Dictionary of users with a list of features they have opted into - """ - - - def parse_user_opt_in_from_text(user_optin_text: str) -> UserOptins: - """ - Parse the user opt-in text into a key value pair of username and the list of features they have opted into - - Users are GitHub usernames with the @ prefix. Each user is also a comma-separated list of features/experiments to enable. - - Example line: "@User1,lf,split_build" - - A "#" prefix indicates the user is opted out of all experiments - - - """ - optins = UserOptins() - for user in user_optin_text.split("\n"): - user = user.strip("\r\n\t -") - if not user or not user.startswith("@"): - # Not a valid user. Skip - continue - - if user: - usr_name = user.split(",")[0].strip("@") - optins[usr_name] = [exp.strip(" ") for exp in user.split(",")[1:]] - - return optins - - - def is_valid_experiment_name(experiment_name: str) -> bool: - """ - Check if the experiment name is valid. - A valid name: - - Contains only alphanumeric characters and the special characters "_" & "-" - - The special characters "_" & "-" shouldn't be the first or last characters - - Cannot contain spaces - """ - - valid_char_regex = r"^[a-zA-Z0-9]([\w-]*[a-zA-Z0-9])?$" - valid = bool(re.match(valid_char_regex, experiment_name)) - - if valid: - return True - - log.error( - f"Invalid experiment name: {experiment_name}. Experiment names should only contain alphanumeric characters, '_', and '-'. They cannot contain spaces, and the special characters '_' and '-' cannot be the first or last characters." - ) - return False - - - def parse_settings_from_text(settings_text: str) -> Settings: - """ - Parse the experiments from the issue body into a list of ExperimentSettings - """ - try: - if settings_text: - # Escape the backtick as well so that we can have the settings in a code block on the GH issue - # for easy reading - # Note: Using ascii for the backtick so that the cat step in _runner-determinator.yml doesn't choke on - # the backtick character in shell commands. - backtick = chr(96) # backtick character - settings_text = settings_text.strip(f"\r\n\t{backtick} ") - settings = load_yaml(settings_text) - - # For now we just load experiments. We can expand this if/when we add more settings - experiments = {} - - for exp_name, exp_settings in settings.get(SETTING_EXPERIMENTS).items(): - if not is_valid_experiment_name(exp_name): - # Exclude invalid experiments from the list. We log an error, but don't raise an exception so that other experiments can still be processed. - continue - - valid_settings = {} - for setting in exp_settings: - if setting not in Experiment._fields: - log.warning( - f"Unexpected setting in experiment: {setting} = {exp_settings[setting]}" - ) - else: - valid_settings[setting] = exp_settings[setting] - - experiments[exp_name] = Experiment(**valid_settings) - return Settings(experiments) - - except Exception: - log.exception("Failed to parse settings") - - return Settings() - - - def parse_settings(rollout_state: str) -> Settings: - """ - Parse settings, if any, from the rollout state. - - If the issue body contains "---" then the text above that is the settings - and the text below is the list of opted in users. - - If it doesn't contain "---" then the settings are empty and the default values are used. - """ - settings_text, _ = extract_settings_user_opt_in_from_text(rollout_state) - return parse_settings_from_text(settings_text) - - - def parse_users(rollout_state: str) -> UserOptins: - """ - Parse users from the rollout state. - - """ - _, users_text = extract_settings_user_opt_in_from_text(rollout_state) - return parse_user_opt_in_from_text(users_text) - - - def is_user_opted_in(user: str, user_optins: UserOptins, experiment_name: str) -> bool: - """ - Check if a user is opted into an experiment - """ - return experiment_name in user_optins.get(user, []) - - - def is_user_opted_out(user: str, user_optins: UserOptins, experiment_name: str) -> bool: - """ - Check if a user explicitly opted out of an experiment - """ - # if the experiment is prefixed with a "-", then it's an opt-out - experiment_optout = "-" + experiment_name - if experiment_optout not in user_optins.get(user, []): - return False - - if is_user_opted_in(user, user_optins, experiment_name): - log.warning( - f"User {user} is opted into experiment {experiment_name}, but also opted out of it. Defaulting to opting out" - ) - - return True - - - def get_runner_prefix( - rollout_state: str, - workflow_requestors: Iterable[str], - branch: str, - eligible_experiments: frozenset[str] = frozenset(), - opt_out_experiments: frozenset[str] = frozenset(), - is_canary: bool = False, - ) -> RunnerPrefixResult: - settings = parse_settings(rollout_state) - user_optins = parse_users(rollout_state) - - fleet_prefix = "" - prefixes = [] - use_arc = False - for experiment_name, experiment_settings in settings.experiments.items(): - if not experiment_settings.all_branches and is_exception_branch(branch): - log.info( - f"Branch {branch} is an exception branch. Not enabling experiment {experiment_name}." - ) - continue - - if opt_out_experiments: - if experiment_name in opt_out_experiments: - opt_out_exp_list = ", ".join(opt_out_experiments) - log.info( - f"Skipping experiment '{experiment_name}', as this workflow has opted-out (opted out experiments are: {opt_out_exp_list})" - ) - continue - - if eligible_experiments: - if experiment_name not in eligible_experiments: - exp_list = ", ".join(eligible_experiments) - log.info( - f"Skipping experiment '{experiment_name}', as it is not in the eligible_experiments list: {exp_list}" - ) - continue - elif not experiment_settings.default: - log.info( - f"Skipping experiment '{experiment_name}', as it is not a default experiment" - ) - continue - - # Is any workflow_requestor opted out to this experiment? - opted_out_users = [ - requestor - for requestor in workflow_requestors - if is_user_opted_out(requestor, user_optins, experiment_name) - ] - - if opted_out_users: - log.info( - f"{', '.join(opted_out_users)} have opted out of experiment {experiment_name}." - ) - continue - - # Is any workflow_requestor opted in to this experiment? - opted_in_users = [ - requestor - for requestor in workflow_requestors - if is_user_opted_in(requestor, user_optins, experiment_name) - ] - - enabled = False - if opted_in_users: - log.info( - f"{', '.join(opted_in_users)} have opted into experiment {experiment_name}." - ) - enabled = True - - elif experiment_settings.rollout_perc: - # If no user is opted in, then we randomly enable the experiment based on the rollout percentage - if random.uniform(0, 100) <= experiment_settings.rollout_perc: - log.info( - f"Based on rollout percentage of {experiment_settings.rollout_perc}%, enabling experiment {experiment_name}." - ) - enabled = True - - if enabled: - label = experiment_name - if experiment_name == ARC_FLEET_EXPERIMENT: - use_arc = True - log.info( - f"ARC experiment enabled. Using ARC runner prefix ({'canary' if is_canary else 'production'})." - ) - elif experiment_name == LF_FLEET_EXPERIMENT: - # We give some special treatment to the "lf" experiment since determines the fleet we use - # - If it's enabled, then we always list it's prefix first - # - If we're in the canary branch, then we append ".c" to the lf prefix - if is_canary: - label += CANARY_FLEET_SUFFIX - fleet_prefix = label - else: - prefixes.append(label) - - # ARC experiment takes precedence: return a fixed label prefix - if use_arc: - arc_prefix = ( - ARC_CANARY_LABEL_PREFIX + ARC_LABEL_PREFIX - if is_canary - else ARC_LABEL_PREFIX - ) - return RunnerPrefixResult(prefix=arc_prefix, use_arc=True) - - if len(prefixes) > 1: - log.error( - f"Only a fleet and one other experiment can be enabled for a job at any time. Enabling {prefixes[0]} and ignoring the rest, which are {', '.join(prefixes[1:])}" - ) - prefixes = prefixes[:1] - - # Fleet always comes first - if fleet_prefix: - prefixes.insert(0, fleet_prefix) - - prefix = ".".join(prefixes) + "." if prefixes else "" - return RunnerPrefixResult(prefix=prefix) - - - def get_rollout_state_from_issue(github_token: str, repo: str, issue_num: int) -> str: - """ - Gets the first comment of the issue, which contains the desired rollout state. - - The default issue we use - https://github.com/pytorch/test-infra/issues/5132 - """ - gh = get_gh_client(github_token) - issue = get_issue(gh, repo, issue_num) - return str(issue.get_comments()[0].body.strip("\n\t ")) - - - def download_json(url: str, headers: dict[str, str], num_retries: int = 3) -> Any: - for _ in range(num_retries): - try: - req = Request(url=url, headers=headers) - content = urlopen(req, timeout=5).read().decode("utf-8") - return json.loads(content) - except Exception as e: - log.warning(f"Could not download {url}: {e}") - - log.warning(f"All {num_retries} retries exhausted, downloading {url} failed") - return {} - - - @cache - def get_pr_info(github_repo: str, github_token: str, pr_number: int) -> dict[str, Any]: - """ - Dynamically get PR information - """ - github_api = f"https://api.github.com/repos/{github_repo}" - headers = { - "Accept": "application/vnd.github.v3+json", - "Authorization": f"token {github_token}", - } - json_response: dict[str, Any] = download_json( - url=f"{github_api}/issues/{pr_number}", - headers=headers, - ) - - if not json_response: - log.warning(f"Failed to get the labels for #{pr_number}") - return {} - - return json_response - - - def get_labels(github_repo: str, github_token: str, pr_number: int) -> set[str]: - """ - Dynamically get the latest list of labels from the pull request - """ - pr_info = get_pr_info(github_repo, github_token, pr_number) - return { - label.get("name") for label in pr_info.get("labels", []) if label.get("name") - } - - - def main() -> None: - args = parse_args() - - runner_label_prefix = DEFAULT_LABEL_PREFIX - - # Check if the PR is opt-out - if args.pr_number: - labels = get_labels(args.github_repo, args.github_token, int(args.pr_number)) - if OPT_OUT_LABEL in labels: - log.info( - f"Opt-out runner determinator because #{args.pr_number} has {OPT_OUT_LABEL} label" - ) - set_github_output(GH_OUTPUT_KEY_LABEL_TYPE, runner_label_prefix) - sys.exit() - - try: - rollout_state = get_rollout_state_from_issue( - args.github_token, args.github_issue_repo, args.github_issue - ) - - username = get_potential_pr_author( - args.github_token, - args.github_repo, - args.github_actor, - args.github_ref_type, - args.github_branch, - ) - - is_canary = args.github_repo == "pytorch/pytorch-canary" - - result = get_runner_prefix( - rollout_state, - (args.github_issue_owner, username), - args.github_branch, - args.eligible_experiments, - args.opt_out_experiments, - is_canary, - ) - runner_label_prefix = result.prefix - set_github_output(GH_OUTPUT_KEY_USE_ARC, str(result.use_arc).lower()) - - except Exception as e: - log.error( - f"Failed to get issue. Defaulting to Meta runners and no experiments. Exception: {e}" - ) - - set_github_output(GH_OUTPUT_KEY_LABEL_TYPE, runner_label_prefix) - - - if __name__ == "__main__": - main() - - EOF - - cat runner_determinator.py + - name: Checkout PyTorch + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - name: Install dependencies run: python3 -m pip install urllib3==1.26.18 PyGithub==2.3.0 @@ -752,7 +76,7 @@ jobs: curr_ref_type="${{ inputs.curr_ref_type }}" echo "Current branch is '$curr_branch'" - python3 runner_determinator.py \ + python3 .github/scripts/runner_determinator.py \ --github-token "$GITHUB_TOKEN" \ --github-issue "$ISSUE_NUMBER" \ --github-branch "$curr_branch" \ diff --git a/.github/workflows/_vllm-benchmark.yml b/.github/workflows/_vllm-benchmark.yml index a97a3190c2ea6..a993d54bac4b9 100644 --- a/.github/workflows/_vllm-benchmark.yml +++ b/.github/workflows/_vllm-benchmark.yml @@ -190,6 +190,7 @@ jobs: --models "${MODELS}" \ --device "${DEVICE_NAME}" \ --include-eager-mode \ + --include-inductor-graph-partition \ --compilation-config "${COMPILATION_CONFIG}" popd diff --git a/.github/workflows/_xpu-test.yml b/.github/workflows/_xpu-test.yml index bd6e0a6fd2563..31d33e6bc9f6a 100644 --- a/.github/workflows/_xpu-test.yml +++ b/.github/workflows/_xpu-test.yml @@ -248,7 +248,7 @@ jobs: -e ZE_AFFINITY_MASK \ -e HUGGING_FACE_HUB_TOKEN \ -e DASHBOARD_TAG \ - --env-file="/tmp/github_env_${GITHUB_RUN_ID}" \ + --env-file="${RUNNER_TEMP}/github_env_${GITHUB_RUN_ID}" \ --ulimit stack=10485760:83886080 \ --ulimit core=0 \ --security-opt seccomp=unconfined \ diff --git a/.github/workflows/attention_op_microbenchmark.yml b/.github/workflows/attention_op_microbenchmark.yml index cd04a48223ce1..085f4cda787e8 100644 --- a/.github/workflows/attention_op_microbenchmark.yml +++ b/.github/workflows/attention_op_microbenchmark.yml @@ -16,6 +16,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: attn-microbenchmark-build: diff --git a/.github/workflows/b200-distributed.yml b/.github/workflows/b200-distributed.yml index e52c7a4b5f5c5..6ce72fd11c442 100644 --- a/.github/workflows/b200-distributed.yml +++ b/.github/workflows/b200-distributed.yml @@ -18,6 +18,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: diff --git a/.github/workflows/b200-symm-mem.yml b/.github/workflows/b200-symm-mem.yml index 62367b61b07b9..256409ace2d6c 100644 --- a/.github/workflows/b200-symm-mem.yml +++ b/.github/workflows/b200-symm-mem.yml @@ -18,6 +18,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: diff --git a/.github/workflows/build-almalinux-images.yml b/.github/workflows/build-almalinux-images.yml index f459c307118fe..2a2b2c8f0caf0 100644 --- a/.github/workflows/build-almalinux-images.yml +++ b/.github/workflows/build-almalinux-images.yml @@ -36,7 +36,7 @@ jobs: runs-on: linux.9xlarge.ephemeral strategy: matrix: - tag: ["cuda12.6", "cuda12.8", "cuda12.9", "cuda13.0", "cuda13.2", "rocm7.0", "rocm7.1", "rocm7.2", "cpu"] + tag: ["cuda12.6", "cuda12.8", "cuda13.0", "cuda13.2", "rocm7.0", "rocm7.1", "rocm7.2", "cpu"] steps: - name: Build docker image uses: pytorch/pytorch/.github/actions/binary-docker-build@main diff --git a/.github/workflows/build-magma-windows.yml b/.github/workflows/build-magma-windows.yml index b7d293a5cec11..39a705338f189 100644 --- a/.github/workflows/build-magma-windows.yml +++ b/.github/workflows/build-magma-windows.yml @@ -22,7 +22,7 @@ jobs: runs-on: windows-2022 strategy: matrix: - cuda_version: ["130", "129", "128", "126"] + cuda_version: ["132", "130", "129", "128", "126"] config: ["Release", "Debug"] env: CUDA_VERSION: ${{ matrix.cuda_version }} diff --git a/.github/workflows/build-manywheel-images.yml b/.github/workflows/build-manywheel-images.yml index 2ada12e59a584..d25dd4a14b1db 100644 --- a/.github/workflows/build-manywheel-images.yml +++ b/.github/workflows/build-manywheel-images.yml @@ -49,13 +49,11 @@ jobs: { name: "manylinux2_28-builder", tag: "cuda13.2", runner: "linux.9xlarge.ephemeral" }, { name: "manylinux2_28-builder", tag: "cuda13.0", runner: "linux.9xlarge.ephemeral" }, { name: "manylinux2_28-builder", tag: "cuda12.8", runner: "linux.9xlarge.ephemeral" }, - { name: "manylinux2_28-builder", tag: "cuda12.9", runner: "linux.9xlarge.ephemeral" }, { name: "manylinux2_28-builder", tag: "cuda12.6", runner: "linux.9xlarge.ephemeral" }, - { name: "manylinuxaarch64-builder", tag: "cuda13.2", runner: "linux.arm64.2xlarge.ephemeral" }, - { name: "manylinuxaarch64-builder", tag: "cuda13.0", runner: "linux.arm64.2xlarge.ephemeral" }, - { name: "manylinuxaarch64-builder", tag: "cuda12.9", runner: "linux.arm64.2xlarge.ephemeral" }, - { name: "manylinuxaarch64-builder", tag: "cuda12.8", runner: "linux.arm64.2xlarge.ephemeral" }, - { name: "manylinuxaarch64-builder", tag: "cuda12.6", runner: "linux.arm64.2xlarge.ephemeral" }, + { name: "manylinuxaarch64-builder", tag: "cuda13.2", runner: "linux.arm64.m7g.4xlarge.ephemeral" }, + { name: "manylinuxaarch64-builder", tag: "cuda13.0", runner: "linux.arm64.m7g.4xlarge.ephemeral" }, + { name: "manylinuxaarch64-builder", tag: "cuda12.8", runner: "linux.arm64.m7g.4xlarge.ephemeral" }, + { name: "manylinuxaarch64-builder", tag: "cuda12.6", runner: "linux.arm64.m7g.4xlarge.ephemeral" }, { name: "manylinux2_28-builder", tag: "rocm7.0", runner: "linux.9xlarge.ephemeral" }, { name: "manylinux2_28-builder", tag: "rocm7.1", runner: "linux.9xlarge.ephemeral" }, { name: "manylinux2_28-builder", tag: "rocm7.2", runner: "linux.9xlarge.ephemeral" }, diff --git a/.github/workflows/build-triton-wheel.yml b/.github/workflows/build-triton-wheel.yml index 49534b2c81bbc..17e5639b01598 100644 --- a/.github/workflows/build-triton-wheel.yml +++ b/.github/workflows/build-triton-wheel.yml @@ -78,14 +78,11 @@ jobs: with: github-secret: ${{ secrets.GITHUB_TOKEN }} - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main + - name: Setup Linux + uses: pytorch/pytorch/.github/actions/setup-linux@main with: submodules: false - - name: Setup Linux - uses: ./.github/actions/setup-linux - - name: Login to ECR uses: ./.github/actions/ecr-login @@ -165,7 +162,7 @@ jobs: docker exec -t "${container_name}" bash -c "${PYTHON_EXECUTABLE} /pytorch/.github/scripts/build_triton_wheel.py --device=$BUILD_DEVICE $RELEASE $WITH_CLANG_LDD" if [[ ("${{ matrix.device }}" == "cuda" || "${{ matrix.device }}" == "xpu") ]]; then - docker exec -t "${container_name}" bash -c "auditwheel repair --plat ${PLATFORM} //artifacts/*.whl" + docker exec -t "${container_name}" bash -c "auditwheel repair --plat ${PLATFORM} --exclude libtriton.so //artifacts/*.whl" else docker exec -t "${container_name}" bash -c "mkdir //artifacts/wheelhouse" docker exec -t "${container_name}" bash -c "mv //artifacts/*.whl //artifacts/wheelhouse/" diff --git a/.github/workflows/build-vllm-wheel.yml b/.github/workflows/build-vllm-wheel.yml index 9865ec6acec81..bca679f323070 100644 --- a/.github/workflows/build-vllm-wheel.yml +++ b/.github/workflows/build-vllm-wheel.yml @@ -69,14 +69,11 @@ jobs: with: github-secret: ${{ secrets.GITHUB_TOKEN }} - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main + - name: Setup Linux + uses: pytorch/pytorch/.github/actions/setup-linux@main with: submodules: false - - name: Setup Linux - uses: ./.github/actions/setup-linux - - name: Login to ECR uses: ./.github/actions/ecr-login diff --git a/.github/workflows/claude-autorevert-advisor.yml b/.github/workflows/claude-autorevert-advisor.yml index 8bbc04a86bf51..19ce91a949abf 100644 --- a/.github/workflows/claude-autorevert-advisor.yml +++ b/.github/workflows/claude-autorevert-advisor.yml @@ -38,15 +38,16 @@ jobs: git fetch origin ${{ inputs.suspect_commit }} --depth=32 || true - name: Configure AWS credentials via OIDC - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_claude_code aws-region: us-east-1 - name: Run AI Advisor id: claude - uses: anthropics/claude-code-action@v1 + uses: anthropics/claude-code-action@6e2bd52842c65e914eba5c8badd17560bd26b5de # v1.0.89 with: + allowed_bots: "pytorch-auto-revert[bot]" use_bedrock: "true" claude_args: | --model global.anthropic.claude-opus-4-6-v1 @@ -127,6 +128,46 @@ jobs: path: /tmp/verdict/verdict.json retention-days: 30 + - name: Upload verdict to S3 for ClickHouse ingestion + if: always() && steps.claude.outputs.structured_output != '' + env: + VERDICT_JSON: ${{ steps.claude.outputs.structured_output }} + SIGNAL_PATTERN: ${{ inputs.signal_pattern }} + run: | + # Build enriched verdict JSON with signal metadata for CH ingestion + jq -n \ + --arg repo "${{ github.repository }}" \ + --argjson run_id "${{ github.run_id }}" \ + --argjson run_attempt "${{ github.run_attempt }}" \ + --arg timestamp "$(date -u +%Y-%m-%dT%H:%M:%S.000)" \ + --arg suspect_commit "${{ inputs.suspect_commit }}" \ + --argjson pr_number "${{ inputs.pr_number }}" \ + --arg signal_key "$(echo "$SIGNAL_PATTERN" | jq -r '.signal_key // ""')" \ + --arg signal_source "$(echo "$SIGNAL_PATTERN" | jq -r '.signal_source // ""')" \ + --arg workflow_name "$(echo "$SIGNAL_PATTERN" | jq -r '.workflow_name // ""')" \ + --arg verdict "$(echo "$VERDICT_JSON" | jq -r '.verdict // ""')" \ + --argjson confidence "$(echo "$VERDICT_JSON" | jq -r '.confidence // 0')" \ + --arg summary "$(echo "$VERDICT_JSON" | jq -r '.summary // ""')" \ + --arg causal_reasoning "$(echo "$VERDICT_JSON" | jq -r '.causal_reasoning // ""')" \ + '{ + repo: $repo, + run_id: $run_id, + run_attempt: $run_attempt, + timestamp: $timestamp, + suspect_commit: $suspect_commit, + pr_number: $pr_number, + signal_key: $signal_key, + signal_source: $signal_source, + workflow_name: $workflow_name, + verdict: $verdict, + confidence: $confidence, + summary: $summary, + causal_reasoning: $causal_reasoning + }' > /tmp/advisor_verdict.json + + aws s3 cp /tmp/advisor_verdict.json \ + "s3://ossci-raw-job-status/autorevert_advisor_verdicts/${{ github.repository }}/${{ github.run_id }}_${{ github.run_attempt }}.json" + - name: Upload usage metrics if: always() uses: pytorch/test-infra/.github/actions/upload-claude-usage@main diff --git a/.github/workflows/claude-code.yml b/.github/workflows/claude-code.yml index 114c480a17e3b..e09b63a1eac89 100644 --- a/.github/workflows/claude-code.yml +++ b/.github/workflows/claude-code.yml @@ -17,3 +17,6 @@ jobs: secrets: inherit with: additional_claude_args: '--allowedTools Skill' + append_system_prompt: | + When asked to review a PR, always use the /pr-review skill first. + It contains PyTorch-specific review guidelines, output format, and critical checks. diff --git a/.github/workflows/claude-issue-triage-run.yml b/.github/workflows/claude-issue-triage-run.yml index 8794a68314221..0173b7fa3b5ef 100644 --- a/.github/workflows/claude-issue-triage-run.yml +++ b/.github/workflows/claude-issue-triage-run.yml @@ -60,22 +60,24 @@ jobs: "ghcr.io/github/github-mcp-server:v0.30.1" ], "env": { - "GITHUB_PERSONAL_ACCESS_TOKEN": "${{ secrets.GITHUB_TOKEN }}" + "GITHUB_PERSONAL_ACCESS_TOKEN": "${GITHUB_TOKEN}" } } } } EOF + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - name: Configure AWS credentials via OIDC - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_claude_code aws-region: us-east-1 - name: Run Issue Triage timeout-minutes: 5 - uses: anthropics/claude-code-action@v1 + uses: anthropics/claude-code-action@6e2bd52842c65e914eba5c8badd17560bd26b5de # v1.0.89 env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} TRIAGE_HOOK_DEBUG_LOG: /tmp/triage_hooks.log diff --git a/.github/workflows/docker-builds.yml b/.github/workflows/docker-builds.yml index 3345c92f65faf..b71ed9df513c8 100644 --- a/.github/workflows/docker-builds.yml +++ b/.github/workflows/docker-builds.yml @@ -53,18 +53,17 @@ jobs: pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11, pytorch-linux-jammy-cuda13.0-cudnn9-py3.12-gcc11-vllm, pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks, - pytorch-linux-jammy-py3.10-clang15, - pytorch-linux-jammy-py3.11-clang15, - pytorch-linux-jammy-py3.12-clang15, - pytorch-linux-jammy-py3.13-clang15, - pytorch-linux-jammy-py3.14-clang15, - pytorch-linux-jammy-py3.14t-clang15, + pytorch-linux-jammy-py3.10-clang18, + pytorch-linux-jammy-py3.11-clang18, + pytorch-linux-jammy-py3.12-clang18, + pytorch-linux-jammy-py3.13-clang18, + pytorch-linux-jammy-py3.14-clang18, + pytorch-linux-jammy-py3.14t-clang18, pytorch-linux-jammy-rocm-n-py3, pytorch-linux-noble-rocm-n-py3, pytorch-linux-noble-rocm-nightly-py3, pytorch-linux-jammy-rocm-n-py3-benchmarks, - pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang15, - pytorch-linux-jammy-py3.10-gcc11, + pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang18, pytorch-linux-jammy-py3-gcc11-inductor-benchmarks, pytorch-linux-jammy-py3.12-halide, pytorch-linux-jammy-py3.12-pallas, @@ -74,12 +73,10 @@ jobs: pytorch-linux-noble-xpu-n-py3, pytorch-linux-noble-xpu-n-py3-client, pytorch-linux-noble-xpu-n-py3-inductor-benchmarks, - pytorch-linux-jammy-py3-clang18-asan, - pytorch-linux-jammy-py3-clang15-onnx, pytorch-linux-jammy-linter, pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-linter, # TODO: Re-enable me when docker pin update happens - # pytorch-linux-jammy-py3-clang15-executorch, + # pytorch-linux-jammy-py3-clang18-executorch, pytorch-linux-jammy-py3.12-triton-cpu, pytorch-linux-noble-riscv64-py3.12-gcc14 ] @@ -100,13 +97,8 @@ jobs: sudo rm -rf "${GITHUB_WORKSPACE}" mkdir "${GITHUB_WORKSPACE}" - # [see note: pytorch repo ref] - # deep clone (fetch-depth 0) required for git merge-base - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - - name: Setup Linux - uses: ./.github/actions/setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main - name: Login to ECR uses: ./.github/actions/ecr-login @@ -159,6 +151,9 @@ jobs: echo $GHCR_PAT | docker login ghcr.io -u pytorch --password-stdin docker tag "${ECR_DOCKER_IMAGE}" "${ghcr_image}:${tag}" docker push "${ghcr_image}:${tag}" + # Also push a tag without the hash for easier reference + docker tag "${ECR_DOCKER_IMAGE}" "${ghcr_image}:${{ matrix.docker-image-name }}" + docker push "${ghcr_image}:${{ matrix.docker-image-name }}" - name: Chown workspace uses: ./.github/actions/chown-workspace diff --git a/.github/workflows/docker-cache-rocm.yml b/.github/workflows/docker-cache-rocm.yml index 9d19bda10e2a7..2462d8aedd124 100644 --- a/.github/workflows/docker-cache-rocm.yml +++ b/.github/workflows/docker-cache-rocm.yml @@ -95,7 +95,7 @@ jobs: docker_image_tag=${{ matrix.docker-image }} docker_image_tag="${docker_image_tag#*:}" # Remove everything before and including first ":" docker_image_tag="${docker_image_tag%-*}" # Remove everything after and including last "-" - ref_name=${{ github.event.workflow_run.head_branch || github.event.inputs.branch }} + ref_name=${WORKFLOW_RUN_HEAD_BRANCH} if [[ $ref_name =~ "release/" ]]; then ref_suffix="release" elif [[ $ref_name == "main" ]]; then @@ -107,3 +107,6 @@ jobs: # mv is atomic operation, so we use intermediate tar.tmp file to prevent read-write contention docker save -o ~/pytorch-data/docker/${docker_image_tag}.tar.tmp ${{ matrix.docker-image }} mv ~/pytorch-data/docker/${docker_image_tag}.tar.tmp ~/pytorch-data/docker/${docker_image_tag}_${ref_suffix}.tar + + env: + WORKFLOW_RUN_HEAD_BRANCH: ${{ github.event.workflow_run.head_branch || github.event.inputs.branch }} diff --git a/.github/workflows/docker-release.yml b/.github/workflows/docker-release.yml index 577a8acb5203f..01a60a213228f 100644 --- a/.github/workflows/docker-release.yml +++ b/.github/workflows/docker-release.yml @@ -86,15 +86,9 @@ jobs: uses: pytorch/test-infra/.github/actions/setup-ssh@main with: github-secret: ${{ secrets.GITHUB_TOKEN }} - # [see note: pytorch repo ref] - # deep clone (fetch-depth 0) required for git merge-base - - name: Checkout PyTorch - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - fetch-depth: 0 - submodules: 'recursive' + - name: Setup Linux - uses: ./.github/actions/setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main - name: Login to ECR uses: ./.github/actions/ecr-login diff --git a/.github/workflows/dtensor.yml b/.github/workflows/dtensor.yml index b56012a885a1b..b45c12b0df10b 100644 --- a/.github/workflows/dtensor.yml +++ b/.github/workflows/dtensor.yml @@ -18,6 +18,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: diff --git a/.github/workflows/dynamo-unittest.yml b/.github/workflows/dynamo-unittest.yml index f7eea350b5644..d8432abc4338c 100644 --- a/.github/workflows/dynamo-unittest.yml +++ b/.github/workflows/dynamo-unittest.yml @@ -17,6 +17,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: @@ -39,8 +40,8 @@ jobs: python-version: ['3.11', '3.12', '3.13'] with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-py${{ matrix.python-version }}-clang15 - docker-image-name: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang15 + build-environment: linux-jammy-py${{ matrix.python-version }}-clang18 + docker-image-name: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang18 test-matrix: | { include: [ { config: "dynamo_core", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" }, @@ -58,8 +59,8 @@ jobs: matrix: python-version: ['3.11', '3.12', '3.13'] with: - build-environment: linux-jammy-py${{ matrix.python-version }}-clang15 - docker-image: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang15 + build-environment: linux-jammy-py${{ matrix.python-version }}-clang18 + docker-image: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang18 test-matrix: | { include: [ { config: "dynamo_core", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" }, diff --git a/.github/workflows/generated-linux-aarch64-binary-manywheel-nightly.yml b/.github/workflows/generated-linux-aarch64-binary-manywheel-nightly.yml index 20c3a18754bba..673e15ac2fe27 100644 --- a/.github/workflows/generated-linux-aarch64-binary-manywheel-nightly.yml +++ b/.github/workflows/generated-linux-aarch64-binary-manywheel-nightly.yml @@ -190,154 +190,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_10-cuda-aarch64-12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.10" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_10-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_10-cuda-aarch64-12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_10-cuda-aarch64-12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.10" - build_name: manywheel-py3_10-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_10-cuda-aarch64-12_8-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_10-cuda-aarch64-12_8-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.10" - build_name: manywheel-py3_10-cuda-aarch64-12_8 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_10-cuda-aarch64-12_9-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.10" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_10-cuda-aarch64-12_9 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_10-cuda-aarch64-12_9-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_10-cuda-aarch64-12_9-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.10" - build_name: manywheel-py3_10-cuda-aarch64-12_9 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_10-cuda-aarch64-12_9-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_10-cuda-aarch64-12_9-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.10" - build_name: manywheel-py3_10-cuda-aarch64-12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_10-cuda-aarch64-13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml @@ -358,7 +210,7 @@ jobs: ALPINE_IMAGE: "arm64v8/alpine" build_name: manywheel-py3_10-cuda-aarch64-13_0 build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' timeout-minutes: 420 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -629,154 +481,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_11-cuda-aarch64-12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.11" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_11-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_11-cuda-aarch64-12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_11-cuda-aarch64-12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.11" - build_name: manywheel-py3_11-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_11-cuda-aarch64-12_8-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_11-cuda-aarch64-12_8-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.11" - build_name: manywheel-py3_11-cuda-aarch64-12_8 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_11-cuda-aarch64-12_9-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.11" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_11-cuda-aarch64-12_9 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_11-cuda-aarch64-12_9-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_11-cuda-aarch64-12_9-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.11" - build_name: manywheel-py3_11-cuda-aarch64-12_9 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_11-cuda-aarch64-12_9-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_11-cuda-aarch64-12_9-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.11" - build_name: manywheel-py3_11-cuda-aarch64-12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_11-cuda-aarch64-13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml @@ -797,7 +501,7 @@ jobs: ALPINE_IMAGE: "arm64v8/alpine" build_name: manywheel-py3_11-cuda-aarch64-13_0 build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' timeout-minutes: 420 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -1068,7 +772,7 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_12-cuda-aarch64-12_8-build: + manywheel-py3_12-cuda-aarch64-13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml needs: get-label-type @@ -1077,26 +781,26 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0-aarch64" GPU_ARCH_TYPE: cuda-aarch64 DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DOCKER_IMAGE_TAG_PREFIX: cuda13.0 DESIRED_PYTHON: "3.12" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.arm64.r7g.12xlarge.memory ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_12-cuda-aarch64-12_8 + build_name: manywheel-py3_12-cuda-aarch64-13_0 build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' timeout-minutes: 420 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda-aarch64-12_8-test: # Testing + manywheel-py3_12-cuda-aarch64-13_0-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - manywheel-py3_12-cuda-aarch64-12_8-build + - manywheel-py3_12-cuda-aarch64-13_0-build - get-label-type uses: ./.github/workflows/_binary-test-linux.yml with: @@ -1104,37 +808,37 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0-aarch64" GPU_ARCH_TYPE: cuda-aarch64 DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DOCKER_IMAGE_TAG_PREFIX: cuda13.0 DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda-aarch64-12_8 + build_name: manywheel-py3_12-cuda-aarch64-13_0 build_environment: linux-aarch64-binary-manywheel runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.arm64.2xlarge ALPINE_IMAGE: "arm64v8/alpine" secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda-aarch64-12_8-upload: # Uploading + manywheel-py3_12-cuda-aarch64-13_0-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: manywheel-py3_12-cuda-aarch64-12_8-build + needs: manywheel-py3_12-cuda-aarch64-13_0-build with: PYTORCH_ROOT: /pytorch PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0-aarch64" GPU_ARCH_TYPE: cuda-aarch64 DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DOCKER_IMAGE_TAG_PREFIX: cuda13.0 DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda-aarch64-12_8 + build_name: manywheel-py3_12-cuda-aarch64-13_0 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -1142,7 +846,7 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_12-cuda-aarch64-12_9-build: + manywheel-py3_12-cuda-aarch64-13_2-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml needs: get-label-type @@ -1151,26 +855,26 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2-aarch64" GPU_ARCH_TYPE: cuda-aarch64 DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda13.2 DESIRED_PYTHON: "3.12" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.arm64.r7g.12xlarge.memory ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_12-cuda-aarch64-12_9 + build_name: manywheel-py3_12-cuda-aarch64-13_2 build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.0; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' timeout-minutes: 420 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda-aarch64-12_9-test: # Testing + manywheel-py3_12-cuda-aarch64-13_2-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - manywheel-py3_12-cuda-aarch64-12_9-build + - manywheel-py3_12-cuda-aarch64-13_2-build - get-label-type uses: ./.github/workflows/_binary-test-linux.yml with: @@ -1178,173 +882,25 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2-aarch64" GPU_ARCH_TYPE: cuda-aarch64 DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda13.2 DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda-aarch64-12_9 + build_name: manywheel-py3_12-cuda-aarch64-13_2 build_environment: linux-aarch64-binary-manywheel runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.arm64.2xlarge ALPINE_IMAGE: "arm64v8/alpine" secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda-aarch64-12_9-upload: # Uploading + manywheel-py3_12-cuda-aarch64-13_2-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: manywheel-py3_12-cuda-aarch64-12_9-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda-aarch64-12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_12-cuda-aarch64-13_0-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.0 - DESIRED_PYTHON: "3.12" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_12-cuda-aarch64-13_0 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_12-cuda-aarch64-13_0-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_12-cuda-aarch64-13_0-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.0 - DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda-aarch64-13_0 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda-aarch64-13_0-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_12-cuda-aarch64-13_0-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.0 - DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda-aarch64-13_0 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_12-cuda-aarch64-13_2-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu132 - GPU_ARCH_VERSION: "13.2-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.2 - DESIRED_PYTHON: "3.12" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_12-cuda-aarch64-13_2 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.0; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_12-cuda-aarch64-13_2-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_12-cuda-aarch64-13_2-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu132 - GPU_ARCH_VERSION: "13.2-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.2 - DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda-aarch64-13_2 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda-aarch64-13_2-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_12-cuda-aarch64-13_2-build + needs: manywheel-py3_12-cuda-aarch64-13_2-build with: PYTORCH_ROOT: /pytorch PACKAGE_TYPE: manywheel @@ -1507,154 +1063,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_13-cuda-aarch64-12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.13" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_13-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_13-cuda-aarch64-12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_13-cuda-aarch64-12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.13" - build_name: manywheel-py3_13-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13-cuda-aarch64-12_8-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_13-cuda-aarch64-12_8-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.13" - build_name: manywheel-py3_13-cuda-aarch64-12_8 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_13-cuda-aarch64-12_9-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.13" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_13-cuda-aarch64-12_9 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_13-cuda-aarch64-12_9-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_13-cuda-aarch64-12_9-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.13" - build_name: manywheel-py3_13-cuda-aarch64-12_9 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13-cuda-aarch64-12_9-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_13-cuda-aarch64-12_9-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.13" - build_name: manywheel-py3_13-cuda-aarch64-12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_13-cuda-aarch64-13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml @@ -1675,7 +1083,7 @@ jobs: ALPINE_IMAGE: "arm64v8/alpine" build_name: manywheel-py3_13-cuda-aarch64-13_0 build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' timeout-minutes: 420 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -1836,183 +1244,35 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cpu - GPU_ARCH_TYPE: cpu-aarch64 - DOCKER_IMAGE: manylinux2_28_aarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cpu-aarch64 - DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cpu-aarch64 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13t-cpu-aarch64-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_13t-cpu-aarch64-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cpu - GPU_ARCH_TYPE: cpu-aarch64 - DOCKER_IMAGE: manylinux2_28_aarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cpu-aarch64 - DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cpu-aarch64 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_13t-cuda-aarch64-12_6-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu126 - GPU_ARCH_VERSION: "12.6-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.6 - DESIRED_PYTHON: "3.13t" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_13t-cuda-aarch64-12_6 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_13t-cuda-aarch64-12_6-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_13t-cuda-aarch64-12_6-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu126 - GPU_ARCH_VERSION: "12.6-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.6 - DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cuda-aarch64-12_6 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13t-cuda-aarch64-12_6-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_13t-cuda-aarch64-12_6-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu126 - GPU_ARCH_VERSION: "12.6-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.6 - DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cuda-aarch64-12_6 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_13t-cuda-aarch64-12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.13t" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_13t-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_13t-cuda-aarch64-12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_13t-cuda-aarch64-12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DESIRED_CUDA: cpu + GPU_ARCH_TYPE: cpu-aarch64 + DOCKER_IMAGE: manylinux2_28_aarch64-builder + DOCKER_IMAGE_TAG_PREFIX: cpu-aarch64 DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cuda-aarch64-12_8 + build_name: manywheel-py3_13t-cpu-aarch64 build_environment: linux-aarch64-binary-manywheel runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.arm64.2xlarge ALPINE_IMAGE: "arm64v8/alpine" secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13t-cuda-aarch64-12_8-upload: # Uploading + manywheel-py3_13t-cpu-aarch64-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: manywheel-py3_13t-cuda-aarch64-12_8-build + needs: manywheel-py3_13t-cpu-aarch64-test with: PYTORCH_ROOT: /pytorch PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DESIRED_CUDA: cpu + GPU_ARCH_TYPE: cpu-aarch64 + DOCKER_IMAGE: manylinux2_28_aarch64-builder + DOCKER_IMAGE_TAG_PREFIX: cpu-aarch64 DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cuda-aarch64-12_8 + build_name: manywheel-py3_13t-cpu-aarch64 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -2020,7 +1280,7 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_13t-cuda-aarch64-12_9-build: + manywheel-py3_13t-cuda-aarch64-12_6-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml needs: get-label-type @@ -2029,26 +1289,26 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" + DESIRED_CUDA: cu126 + GPU_ARCH_VERSION: "12.6-aarch64" GPU_ARCH_TYPE: cuda-aarch64 DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda12.6 DESIRED_PYTHON: "3.13t" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.arm64.r7g.12xlarge.memory ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_13t-cuda-aarch64-12_9 + build_name: manywheel-py3_13t-cuda-aarch64-12_6 build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' timeout-minutes: 420 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13t-cuda-aarch64-12_9-test: # Testing + manywheel-py3_13t-cuda-aarch64-12_6-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - manywheel-py3_13t-cuda-aarch64-12_9-build + - manywheel-py3_13t-cuda-aarch64-12_6-build - get-label-type uses: ./.github/workflows/_binary-test-linux.yml with: @@ -2056,37 +1316,37 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" + DESIRED_CUDA: cu126 + GPU_ARCH_VERSION: "12.6-aarch64" GPU_ARCH_TYPE: cuda-aarch64 DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda12.6 DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cuda-aarch64-12_9 + build_name: manywheel-py3_13t-cuda-aarch64-12_6 build_environment: linux-aarch64-binary-manywheel runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.arm64.2xlarge ALPINE_IMAGE: "arm64v8/alpine" secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13t-cuda-aarch64-12_9-upload: # Uploading + manywheel-py3_13t-cuda-aarch64-12_6-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: manywheel-py3_13t-cuda-aarch64-12_9-build + needs: manywheel-py3_13t-cuda-aarch64-12_6-build with: PYTORCH_ROOT: /pytorch PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" + DESIRED_CUDA: cu126 + GPU_ARCH_VERSION: "12.6-aarch64" GPU_ARCH_TYPE: cuda-aarch64 DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda12.6 DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cuda-aarch64-12_9 + build_name: manywheel-py3_13t-cuda-aarch64-12_6 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -2114,7 +1374,7 @@ jobs: ALPINE_IMAGE: "arm64v8/alpine" build_name: manywheel-py3_13t-cuda-aarch64-13_0 build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' timeout-minutes: 420 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -2385,154 +1645,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_14-cuda-aarch64-12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.14" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_14-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_14-cuda-aarch64-12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_14-cuda-aarch64-12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14-cuda-aarch64-12_8-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_14-cuda-aarch64-12_8-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cuda-aarch64-12_8 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_14-cuda-aarch64-12_9-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.14" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_14-cuda-aarch64-12_9 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_14-cuda-aarch64-12_9-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_14-cuda-aarch64-12_9-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cuda-aarch64-12_9 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14-cuda-aarch64-12_9-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_14-cuda-aarch64-12_9-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cuda-aarch64-12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_14-cuda-aarch64-13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml @@ -2553,7 +1665,7 @@ jobs: ALPINE_IMAGE: "arm64v8/alpine" build_name: manywheel-py3_14-cuda-aarch64-13_0 build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' timeout-minutes: 420 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -2824,154 +1936,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_14t-cuda-aarch64-12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.14t" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_14t-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_14t-cuda-aarch64-12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_14t-cuda-aarch64-12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.14t" - build_name: manywheel-py3_14t-cuda-aarch64-12_8 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14t-cuda-aarch64-12_8-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_14t-cuda-aarch64-12_8-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.14t" - build_name: manywheel-py3_14t-cuda-aarch64-12_8 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_14t-cuda-aarch64-12_9-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.14t" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.r7g.12xlarge.memory - ALPINE_IMAGE: "arm64v8/alpine" - build_name: manywheel-py3_14t-cuda-aarch64-12_9 - build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - timeout-minutes: 420 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_14t-cuda-aarch64-12_9-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_14t-cuda-aarch64-12_9-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.14t" - build_name: manywheel-py3_14t-cuda-aarch64-12_9 - build_environment: linux-aarch64-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.arm64.2xlarge - ALPINE_IMAGE: "arm64v8/alpine" - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14t-cuda-aarch64-12_9-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_14t-cuda-aarch64-12_9-build - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9-aarch64" - GPU_ARCH_TYPE: cuda-aarch64 - DOCKER_IMAGE: manylinuxaarch64-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.14t" - build_name: manywheel-py3_14t-cuda-aarch64-12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_14t-cuda-aarch64-13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml @@ -2992,7 +1956,7 @@ jobs: ALPINE_IMAGE: "arm64v8/alpine" build_name: manywheel-py3_14t-cuda-aarch64-13_0 build_environment: linux-aarch64-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' timeout-minutes: 420 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/generated-linux-binary-manywheel-nightly.yml b/.github/workflows/generated-linux-binary-manywheel-nightly.yml index 9e71c7223a9ef..adbb922345d10 100644 --- a/.github/workflows/generated-linux-binary-manywheel-nightly.yml +++ b/.github/workflows/generated-linux-binary-manywheel-nightly.yml @@ -184,146 +184,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_10-cuda12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.10" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_10-cuda12_8 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_10-cuda12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_10-cuda12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.10" - build_name: manywheel-py3_10-cuda12_8 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_10-cuda12_8-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_10-cuda12_8-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.10" - build_name: manywheel-py3_10-cuda12_8 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_10-cuda12_9-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.10" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_10-cuda12_9 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_10-cuda12_9-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_10-cuda12_9-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.10" - build_name: manywheel-py3_10-cuda12_9 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_10-cuda12_9-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_10-cuda12_9-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.10" - build_name: manywheel-py3_10-cuda12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_10-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml @@ -342,7 +202,7 @@ jobs: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build_name: manywheel-py3_10-cuda13_0 build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -480,7 +340,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.1 DESIRED_PYTHON: "3.10" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_10-rocm7_1 build_environment: linux-binary-manywheel secrets: @@ -532,7 +392,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -596,7 +456,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.2 DESIRED_PYTHON: "3.10" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_10-rocm7_2 build_environment: linux-binary-manywheel secrets: @@ -648,7 +508,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -936,146 +796,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_11-cuda12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.11" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_11-cuda12_8 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_11-cuda12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_11-cuda12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.11" - build_name: manywheel-py3_11-cuda12_8 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_11-cuda12_8-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_11-cuda12_8-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.11" - build_name: manywheel-py3_11-cuda12_8 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_11-cuda12_9-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.11" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_11-cuda12_9 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_11-cuda12_9-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_11-cuda12_9-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.11" - build_name: manywheel-py3_11-cuda12_9 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_11-cuda12_9-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_11-cuda12_9-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.11" - build_name: manywheel-py3_11-cuda12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_11-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml @@ -1094,7 +814,7 @@ jobs: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build_name: manywheel-py3_11-cuda13_0 build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -1232,7 +952,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.1 DESIRED_PYTHON: "3.11" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_11-rocm7_1 build_environment: linux-binary-manywheel secrets: @@ -1284,7 +1004,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -1348,7 +1068,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.2 DESIRED_PYTHON: "3.11" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_11-rocm7_2 build_environment: linux-binary-manywheel secrets: @@ -1400,7 +1120,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -1688,7 +1408,7 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_12-cuda12_8-build: + manywheel-py3_12-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml needs: get-label-type @@ -1697,23 +1417,23 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DOCKER_IMAGE_TAG_PREFIX: cuda13.0 DESIRED_PYTHON: "3.12" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_12-cuda12_8 + build_name: manywheel-py3_12-cuda13_0 build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda12_8-test: # Testing + manywheel-py3_12-cuda13_0-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - manywheel-py3_12-cuda12_8-build + - manywheel-py3_12-cuda13_0-build - get-label-type uses: ./.github/workflows/_binary-test-linux.yml with: @@ -1721,36 +1441,36 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DOCKER_IMAGE_TAG_PREFIX: cuda13.0 DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda12_8 + build_name: manywheel-py3_12-cuda13_0 build_environment: linux-binary-manywheel runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.g4dn.4xlarge.nvidia.gpu secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda12_8-upload: # Uploading + manywheel-py3_12-cuda13_0-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: manywheel-py3_12-cuda12_8-test + needs: manywheel-py3_12-cuda13_0-test with: PYTORCH_ROOT: /pytorch PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DOCKER_IMAGE_TAG_PREFIX: cuda13.0 DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda12_8 + build_name: manywheel-py3_12-cuda13_0 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -1758,7 +1478,7 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_12-cuda12_9-build: + manywheel-py3_12-cuda13_2-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml needs: get-label-type @@ -1767,23 +1487,23 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda13.2 DESIRED_PYTHON: "3.12" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_12-cuda12_9 + build_name: manywheel-py3_12-cuda13_2 build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.0; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda12_9-test: # Testing + manywheel-py3_12-cuda13_2-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - manywheel-py3_12-cuda12_9-build + - manywheel-py3_12-cuda13_2-build - get-label-type uses: ./.github/workflows/_binary-test-linux.yml with: @@ -1791,36 +1511,36 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda13.2 DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda12_9 + build_name: manywheel-py3_12-cuda13_2 build_environment: linux-binary-manywheel runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.g4dn.4xlarge.nvidia.gpu secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda12_9-upload: # Uploading + manywheel-py3_12-cuda13_2-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: manywheel-py3_12-cuda12_9-test + needs: manywheel-py3_12-cuda13_2-test with: PYTORCH_ROOT: /pytorch PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda13.2 DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda12_9 + build_name: manywheel-py3_12-cuda13_2 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -1828,7 +1548,7 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_12-cuda13_0-build: + manywheel-py3_12-rocm7_1-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml needs: get-label-type @@ -1837,155 +1557,15 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" - GPU_ARCH_TYPE: cuda + DESIRED_CUDA: rocm7.1 + GPU_ARCH_VERSION: "7.1" + GPU_ARCH_TYPE: rocm DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.0 + DOCKER_IMAGE_TAG_PREFIX: rocm7.1 DESIRED_PYTHON: "3.12" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_12-cuda13_0 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_12-cuda13_0-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_12-cuda13_0-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.0 - DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda13_0 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda13_0-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_12-cuda13_0-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.0 - DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda13_0 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_12-cuda13_2-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu132 - GPU_ARCH_VERSION: "13.2" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.2 - DESIRED_PYTHON: "3.12" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_12-cuda13_2 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.2.0; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_12-cuda13_2-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_12-cuda13_2-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu132 - GPU_ARCH_VERSION: "13.2" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.2 - DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda13_2 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_12-cuda13_2-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_12-cuda13_2-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu132 - GPU_ARCH_VERSION: "13.2" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda13.2 - DESIRED_PYTHON: "3.12" - build_name: manywheel-py3_12-cuda13_2 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_12-rocm7_1-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: rocm7.1 - GPU_ARCH_VERSION: "7.1" - GPU_ARCH_TYPE: rocm - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: rocm7.1 - DESIRED_PYTHON: "3.12" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 - build_name: manywheel-py3_12-rocm7_1 + timeout-minutes: 420 + build_name: manywheel-py3_12-rocm7_1 build_environment: linux-binary-manywheel secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -2036,7 +1616,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -2100,7 +1680,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.2 DESIRED_PYTHON: "3.12" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_12-rocm7_2 build_environment: linux-binary-manywheel secrets: @@ -2152,7 +1732,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -2440,146 +2020,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_13-cuda12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.13" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_13-cuda12_8 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_13-cuda12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_13-cuda12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.13" - build_name: manywheel-py3_13-cuda12_8 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13-cuda12_8-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_13-cuda12_8-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.13" - build_name: manywheel-py3_13-cuda12_8 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_13-cuda12_9-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.13" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_13-cuda12_9 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_13-cuda12_9-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_13-cuda12_9-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.13" - build_name: manywheel-py3_13-cuda12_9 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13-cuda12_9-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_13-cuda12_9-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.13" - build_name: manywheel-py3_13-cuda12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_13-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml @@ -2598,7 +2038,7 @@ jobs: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build_name: manywheel-py3_13-cuda13_0 build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -2736,7 +2176,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.1 DESIRED_PYTHON: "3.13" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_13-rocm7_1 build_environment: linux-binary-manywheel secrets: @@ -2788,7 +2228,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -2852,7 +2292,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.2 DESIRED_PYTHON: "3.13" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_13-rocm7_2 build_environment: linux-binary-manywheel secrets: @@ -2904,7 +2344,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -3192,146 +2632,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_13t-cuda12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.13t" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_13t-cuda12_8 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_13t-cuda12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_13t-cuda12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cuda12_8 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13t-cuda12_8-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_13t-cuda12_8-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cuda12_8 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_13t-cuda12_9-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.13t" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_13t-cuda12_9 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_13t-cuda12_9-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_13t-cuda12_9-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cuda12_9 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_13t-cuda12_9-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_13t-cuda12_9-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-cuda12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_13t-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml @@ -3350,7 +2650,7 @@ jobs: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build_name: manywheel-py3_13t-cuda13_0 build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -3488,7 +2788,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.1 DESIRED_PYTHON: "3.13t" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_13t-rocm7_1 build_environment: linux-binary-manywheel secrets: @@ -3540,7 +2840,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -3604,7 +2904,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.2 DESIRED_PYTHON: "3.13t" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_13t-rocm7_2 build_environment: linux-binary-manywheel secrets: @@ -3656,7 +2956,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -3793,150 +3093,14 @@ jobs: with: PYTORCH_ROOT: /pytorch PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: xpu - GPU_ARCH_TYPE: xpu - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: xpu - DESIRED_PYTHON: "3.13t" - build_name: manywheel-py3_13t-xpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_14-cpu-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cpu - GPU_ARCH_TYPE: cpu - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cpu - DESIRED_PYTHON: "3.14" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_14-cpu - build_environment: linux-binary-manywheel - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_14-cpu-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_14-cpu-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cpu - GPU_ARCH_TYPE: cpu - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cpu - DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cpu - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.4xlarge - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14-cpu-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_14-cpu-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cpu - GPU_ARCH_TYPE: cpu - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cpu - DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_14-cuda12_6-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu126 - GPU_ARCH_VERSION: "12.6" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.6 - DESIRED_PYTHON: "3.14" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_14-cuda12_6 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_14-cuda12_6-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_14-cuda12_6-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu126 - GPU_ARCH_VERSION: "12.6" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.6 - DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cuda12_6 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14-cuda12_6-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_14-cuda12_6-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu126 - GPU_ARCH_VERSION: "12.6" - GPU_ARCH_TYPE: cuda + # TODO: This is a legacy variable that we eventually want to get rid of in + # favor of GPU_ARCH_VERSION + DESIRED_CUDA: xpu + GPU_ARCH_TYPE: xpu DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.6 - DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cuda12_6 + DOCKER_IMAGE_TAG_PREFIX: xpu + DESIRED_PYTHON: "3.13t" + build_name: manywheel-py3_13t-xpu secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -3944,7 +3108,7 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_14-cuda12_8-build: + manywheel-py3_14-cpu-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml needs: get-label-type @@ -3953,23 +3117,21 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda + DESIRED_CUDA: cpu + GPU_ARCH_TYPE: cpu DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DOCKER_IMAGE_TAG_PREFIX: cpu DESIRED_PYTHON: "3.14" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_14-cuda12_8 + build_name: manywheel-py3_14-cpu build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14-cuda12_8-test: # Testing + manywheel-py3_14-cpu-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - manywheel-py3_14-cuda12_8-build + - manywheel-py3_14-cpu-build - get-label-type uses: ./.github/workflows/_binary-test-linux.yml with: @@ -3977,36 +3139,34 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda + DESIRED_CUDA: cpu + GPU_ARCH_TYPE: cpu DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DOCKER_IMAGE_TAG_PREFIX: cpu DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cuda12_8 + build_name: manywheel-py3_14-cpu build_environment: linux-binary-manywheel runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu + runs_on: linux.4xlarge secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14-cuda12_8-upload: # Uploading + manywheel-py3_14-cpu-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: manywheel-py3_14-cuda12_8-test + needs: manywheel-py3_14-cpu-test with: PYTORCH_ROOT: /pytorch PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda + DESIRED_CUDA: cpu + GPU_ARCH_TYPE: cpu DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 + DOCKER_IMAGE_TAG_PREFIX: cpu DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cuda12_8 + build_name: manywheel-py3_14-cpu secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -4014,7 +3174,7 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_14-cuda12_9-build: + manywheel-py3_14-cuda12_6-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml needs: get-label-type @@ -4023,23 +3183,23 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" + DESIRED_CUDA: cu126 + GPU_ARCH_VERSION: "12.6" GPU_ARCH_TYPE: cuda DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda12.6 DESIRED_PYTHON: "3.14" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_14-cuda12_9 + build_name: manywheel-py3_14-cuda12_6 build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.6.3; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.10.2.21; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.3; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14-cuda12_9-test: # Testing + manywheel-py3_14-cuda12_6-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - manywheel-py3_14-cuda12_9-build + - manywheel-py3_14-cuda12_6-build - get-label-type uses: ./.github/workflows/_binary-test-linux.yml with: @@ -4047,36 +3207,36 @@ jobs: PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" + DESIRED_CUDA: cu126 + GPU_ARCH_VERSION: "12.6" GPU_ARCH_TYPE: cuda DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda12.6 DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cuda12_9 + build_name: manywheel-py3_14-cuda12_6 build_environment: linux-binary-manywheel runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runs_on: linux.g4dn.4xlarge.nvidia.gpu secrets: github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14-cuda12_9-upload: # Uploading + manywheel-py3_14-cuda12_6-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: manywheel-py3_14-cuda12_9-test + needs: manywheel-py3_14-cuda12_6-test with: PYTORCH_ROOT: /pytorch PACKAGE_TYPE: manywheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" + DESIRED_CUDA: cu126 + GPU_ARCH_VERSION: "12.6" GPU_ARCH_TYPE: cuda DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 + DOCKER_IMAGE_TAG_PREFIX: cuda12.6 DESIRED_PYTHON: "3.14" - build_name: manywheel-py3_14-cuda12_9 + build_name: manywheel-py3_14-cuda12_6 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -4102,7 +3262,7 @@ jobs: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build_name: manywheel-py3_14-cuda13_0 build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -4240,7 +3400,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.1 DESIRED_PYTHON: "3.14" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_14-rocm7_1 build_environment: linux-binary-manywheel secrets: @@ -4292,7 +3452,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -4356,7 +3516,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.2 DESIRED_PYTHON: "3.14" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_14-rocm7_2 build_environment: linux-binary-manywheel secrets: @@ -4408,7 +3568,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -4696,146 +3856,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_14t-cuda12_8-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.14t" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_14t-cuda12_8 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.8.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.7.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_14t-cuda12_8-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_14t-cuda12_8-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.14t" - build_name: manywheel-py3_14t-cuda12_8 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14t-cuda12_8-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_14t-cuda12_8-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.8 - DESIRED_PYTHON: "3.14t" - build_name: manywheel-py3_14t-cuda12_8 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - manywheel-py3_14t-cuda12_9-build: - if: ${{ github.repository_owner == 'pytorch' }} - uses: ./.github/workflows/_binary-build-linux.yml - needs: get-label-type - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.14t" - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build_name: manywheel-py3_14t-cuda12_9 - build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==12.9.1; platform_system == 'Linux' | cuda-bindings>=12.9.4,<13; platform_system == 'Linux' | nvidia-cudnn-cu12==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu12==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu12==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu12==3.4.5; platform_system == 'Linux' - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - - manywheel-py3_14t-cuda12_9-test: # Testing - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_14t-cuda12_9-build - - get-label-type - uses: ./.github/workflows/_binary-test-linux.yml - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.14t" - build_name: manywheel-py3_14t-cuda12_9 - build_environment: linux-binary-manywheel - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - runs_on: linux.g4dn.4xlarge.nvidia.gpu - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - manywheel-py3_14t-cuda12_9-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: manywheel-py3_14t-cuda12_9-test - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: manywheel - # TODO: This is a legacy variable that we eventually want to get rid of in - # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - DOCKER_IMAGE: manylinux2_28-builder - DOCKER_IMAGE_TAG_PREFIX: cuda12.9 - DESIRED_PYTHON: "3.14t" - build_name: manywheel-py3_14t-cuda12_9 - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - manywheel-py3_14t-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} uses: ./.github/workflows/_binary-build-linux.yml @@ -4854,7 +3874,7 @@ jobs: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build_name: manywheel-py3_14t-cuda13_0 build_environment: linux-binary-manywheel - PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cublas,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' secrets: github-token: ${{ secrets.GITHUB_TOKEN }} @@ -4992,7 +4012,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.1 DESIRED_PYTHON: "3.14t" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_14t-rocm7_1 build_environment: linux-binary-manywheel secrets: @@ -5044,7 +4064,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -5108,7 +4128,7 @@ jobs: DOCKER_IMAGE_TAG_PREFIX: rocm7.2 DESIRED_PYTHON: "3.14t" runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - timeout-minutes: 300 + timeout-minutes: 420 build_name: manywheel-py3_14t-rocm7_2 build_environment: linux-binary-manywheel secrets: @@ -5160,7 +4180,7 @@ jobs: - name: configure aws credentials id: aws_creds if: ${{ startsWith(github.event.ref, 'refs/tags/ciflow/') }} - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_s3_and_ecr_read_only aws-region: us-east-1 @@ -5433,128 +4453,6 @@ jobs: R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - libtorch-cuda12_8-shared-with-deps-release-extract: - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_10-cuda12_8-build - - get-label-type - runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" - timeout-minutes: 60 - env: - DESIRED_CUDA: cu128 - LIBTORCH_VARIANT: shared-with-deps - steps: - - name: Checkout PyTorch - uses: actions/checkout@v4 - with: - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} - sparse-checkout: .ci/libtorch/ - show-progress: false - - uses: actions/download-artifact@v4.1.7 - name: Download Wheel Artifact - with: - name: manywheel-py3_10-cuda12_8 - path: "${{ runner.temp }}/wheel_artifact/" - - name: Extract libtorch from wheel - run: | - set -eux - mkdir -p "${{ runner.temp }}/libtorch_output" - python3 .ci/libtorch/extract_libtorch_from_wheel.py \ - --wheel-dir "${{ runner.temp }}/wheel_artifact" \ - --output-dir "${{ runner.temp }}/libtorch_output" \ - --platform linux \ - --desired-cuda "$DESIRED_CUDA" \ - --libtorch-variant "$LIBTORCH_VARIANT" - - uses: actions/upload-artifact@v4.4.0 - if: always() - with: - name: libtorch-cuda12_8-shared-with-deps-release - retention-days: 14 - if-no-files-found: error - path: "${{ runner.temp }}/libtorch_output/" - libtorch-cuda12_8-shared-with-deps-release-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: libtorch-cuda12_8-shared-with-deps-release-extract - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: libtorch - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" - GPU_ARCH_TYPE: cuda - LIBTORCH_CONFIG: release - LIBTORCH_VARIANT: shared-with-deps - build_name: libtorch-cuda12_8-shared-with-deps-release - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - - libtorch-cuda12_9-shared-with-deps-release-extract: - if: ${{ github.repository_owner == 'pytorch' }} - needs: - - manywheel-py3_10-cuda12_9-build - - get-label-type - runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" - timeout-minutes: 60 - env: - DESIRED_CUDA: cu129 - LIBTORCH_VARIANT: shared-with-deps - steps: - - name: Checkout PyTorch - uses: actions/checkout@v4 - with: - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} - sparse-checkout: .ci/libtorch/ - show-progress: false - - uses: actions/download-artifact@v4.1.7 - name: Download Wheel Artifact - with: - name: manywheel-py3_10-cuda12_9 - path: "${{ runner.temp }}/wheel_artifact/" - - name: Extract libtorch from wheel - run: | - set -eux - mkdir -p "${{ runner.temp }}/libtorch_output" - python3 .ci/libtorch/extract_libtorch_from_wheel.py \ - --wheel-dir "${{ runner.temp }}/wheel_artifact" \ - --output-dir "${{ runner.temp }}/libtorch_output" \ - --platform linux \ - --desired-cuda "$DESIRED_CUDA" \ - --libtorch-variant "$LIBTORCH_VARIANT" - - uses: actions/upload-artifact@v4.4.0 - if: always() - with: - name: libtorch-cuda12_9-shared-with-deps-release - retention-days: 14 - if-no-files-found: error - path: "${{ runner.temp }}/libtorch_output/" - libtorch-cuda12_9-shared-with-deps-release-upload: # Uploading - if: ${{ github.repository_owner == 'pytorch' }} - permissions: - id-token: write - contents: read - needs: libtorch-cuda12_9-shared-with-deps-release-extract - with: - PYTORCH_ROOT: /pytorch - PACKAGE_TYPE: libtorch - DESIRED_CUDA: cu129 - GPU_ARCH_VERSION: "12.9" - GPU_ARCH_TYPE: cuda - LIBTORCH_CONFIG: release - LIBTORCH_VARIANT: shared-with-deps - build_name: libtorch-cuda12_9-shared-with-deps-release - secrets: - github-token: ${{ secrets.GITHUB_TOKEN }} - R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} - R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} - R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} - uses: ./.github/workflows/_binary-upload.yml - libtorch-cuda13_0-shared-with-deps-release-extract: if: ${{ github.repository_owner == 'pytorch' }} needs: diff --git a/.github/workflows/generated-macos-arm64-binary-wheel-nightly.yml b/.github/workflows/generated-macos-arm64-binary-wheel-nightly.yml index e336e52c48434..c1d956c5bdc1f 100644 --- a/.github/workflows/generated-macos-arm64-binary-wheel-nightly.yml +++ b/.github/workflows/generated-macos-arm64-binary-wheel-nightly.yml @@ -33,7 +33,7 @@ concurrency: jobs: wheel-py3_10-cpu-build: if: ${{ github.repository_owner == 'pytorch' }} - runs-on: macos-14-xlarge + runs-on: macos-26-xlarge timeout-minutes: 240 env: PYTORCH_ROOT: ${{ github.workspace }} @@ -44,6 +44,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.10" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -144,7 +145,7 @@ jobs: uses: ./.github/workflows/_binary-upload.yml wheel-py3_11-cpu-build: if: ${{ github.repository_owner == 'pytorch' }} - runs-on: macos-14-xlarge + runs-on: macos-26-xlarge timeout-minutes: 240 env: PYTORCH_ROOT: ${{ github.workspace }} @@ -155,6 +156,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.11" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -255,7 +257,7 @@ jobs: uses: ./.github/workflows/_binary-upload.yml wheel-py3_12-cpu-build: if: ${{ github.repository_owner == 'pytorch' }} - runs-on: macos-14-xlarge + runs-on: macos-26-xlarge timeout-minutes: 240 env: PYTORCH_ROOT: ${{ github.workspace }} @@ -266,6 +268,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.12" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -366,7 +369,7 @@ jobs: uses: ./.github/workflows/_binary-upload.yml wheel-py3_13-cpu-build: if: ${{ github.repository_owner == 'pytorch' }} - runs-on: macos-14-xlarge + runs-on: macos-26-xlarge timeout-minutes: 240 env: PYTORCH_ROOT: ${{ github.workspace }} @@ -377,6 +380,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -477,7 +481,7 @@ jobs: uses: ./.github/workflows/_binary-upload.yml wheel-py3_13t-cpu-build: if: ${{ github.repository_owner == 'pytorch' }} - runs-on: macos-14-xlarge + runs-on: macos-26-xlarge timeout-minutes: 240 env: PYTORCH_ROOT: ${{ github.workspace }} @@ -488,6 +492,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13t" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -588,7 +593,7 @@ jobs: uses: ./.github/workflows/_binary-upload.yml wheel-py3_14-cpu-build: if: ${{ github.repository_owner == 'pytorch' }} - runs-on: macos-14-xlarge + runs-on: macos-26-xlarge timeout-minutes: 240 env: PYTORCH_ROOT: ${{ github.workspace }} @@ -599,6 +604,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -699,7 +705,7 @@ jobs: uses: ./.github/workflows/_binary-upload.yml wheel-py3_14t-cpu-build: if: ${{ github.repository_owner == 'pytorch' }} - runs-on: macos-14-xlarge + runs-on: macos-26-xlarge timeout-minutes: 240 env: PYTORCH_ROOT: ${{ github.workspace }} @@ -710,6 +716,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14t" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -811,7 +818,7 @@ jobs: libtorch-cpu-shared-with-deps-release-extract: if: ${{ github.repository_owner == 'pytorch' }} needs: wheel-py3_10-cpu-build - runs-on: macos-14-xlarge + runs-on: macos-26-xlarge timeout-minutes: 60 env: DESIRED_CUDA: cpu diff --git a/.github/workflows/generated-windows-binary-libtorch-debug-nightly.yml b/.github/workflows/generated-windows-binary-libtorch-debug-nightly.yml index cb2de611421fe..1bb00a665713b 100644 --- a/.github/workflows/generated-windows-binary-libtorch-debug-nightly.yml +++ b/.github/workflows/generated-windows-binary-libtorch-debug-nightly.yml @@ -528,7 +528,7 @@ jobs: R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - libtorch-cuda12_8-shared-with-deps-debug-build: + libtorch-cuda13_0-shared-with-deps-debug-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -538,8 +538,8 @@ jobs: PACKAGE_TYPE: libtorch # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 LIBTORCH_CONFIG: debug @@ -622,7 +622,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: libtorch-cuda12_8-shared-with-deps-debug + name: libtorch-cuda13_0-shared-with-deps-debug retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -638,10 +638,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - libtorch-cuda12_8-shared-with-deps-debug-test: # Testing + libtorch-cuda13_0-shared-with-deps-debug-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - libtorch-cuda12_8-shared-with-deps-debug-build + - libtorch-cuda13_0-shared-with-deps-debug-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -650,8 +650,8 @@ jobs: PACKAGE_TYPE: libtorch # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 LIBTORCH_CONFIG: debug @@ -726,7 +726,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: libtorch-cuda12_8-shared-with-deps-debug + name: libtorch-cuda13_0-shared-with-deps-debug path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -747,33 +747,33 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - libtorch-cuda12_8-shared-with-deps-debug-upload: # Uploading + libtorch-cuda13_0-shared-with-deps-debug-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: libtorch-cuda12_8-shared-with-deps-debug-test + needs: libtorch-cuda13_0-shared-with-deps-debug-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: libtorch # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda LIBTORCH_CONFIG: debug LIBTORCH_VARIANT: shared-with-deps # This is a dummy value for libtorch to work correctly with our batch scripts # without this value pip does not get installed for some reason DESIRED_PYTHON: "3.10" - build_name: libtorch-cuda12_8-shared-with-deps-debug + build_name: libtorch-cuda13_0-shared-with-deps-debug secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - libtorch-cuda13_0-shared-with-deps-debug-build: + libtorch-cuda13_2-shared-with-deps-debug-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -783,8 +783,8 @@ jobs: PACKAGE_TYPE: libtorch # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 LIBTORCH_CONFIG: debug @@ -867,7 +867,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: libtorch-cuda13_0-shared-with-deps-debug + name: libtorch-cuda13_2-shared-with-deps-debug retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -883,10 +883,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - libtorch-cuda13_0-shared-with-deps-debug-test: # Testing + libtorch-cuda13_2-shared-with-deps-debug-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - libtorch-cuda13_0-shared-with-deps-debug-build + - libtorch-cuda13_2-shared-with-deps-debug-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -895,8 +895,8 @@ jobs: PACKAGE_TYPE: libtorch # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 LIBTORCH_CONFIG: debug @@ -971,7 +971,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: libtorch-cuda13_0-shared-with-deps-debug + name: libtorch-cuda13_2-shared-with-deps-debug path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -992,26 +992,26 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - libtorch-cuda13_0-shared-with-deps-debug-upload: # Uploading + libtorch-cuda13_2-shared-with-deps-debug-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: libtorch-cuda13_0-shared-with-deps-debug-test + needs: libtorch-cuda13_2-shared-with-deps-debug-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: libtorch # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda LIBTORCH_CONFIG: debug LIBTORCH_VARIANT: shared-with-deps # This is a dummy value for libtorch to work correctly with our batch scripts # without this value pip does not get installed for some reason DESIRED_PYTHON: "3.10" - build_name: libtorch-cuda13_0-shared-with-deps-debug + build_name: libtorch-cuda13_2-shared-with-deps-debug secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} diff --git a/.github/workflows/generated-windows-binary-wheel-nightly.yml b/.github/workflows/generated-windows-binary-wheel-nightly.yml index f92df48b5d86f..7da288111208d 100644 --- a/.github/workflows/generated-windows-binary-wheel-nightly.yml +++ b/.github/workflows/generated-windows-binary-wheel-nightly.yml @@ -56,6 +56,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.10" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -505,7 +506,7 @@ jobs: R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_10-cuda12_8-build: + wheel-py3_10-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -515,8 +516,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.10" @@ -595,7 +596,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_10-cuda12_8 + name: wheel-py3_10-cuda13_0 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -611,10 +612,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_10-cuda12_8-test: # Testing + wheel-py3_10-cuda13_0-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_10-cuda12_8-build + - wheel-py3_10-cuda13_0-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -623,8 +624,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.10" @@ -695,7 +696,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_10-cuda12_8 + name: wheel-py3_10-cuda13_0 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -716,29 +717,29 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_10-cuda12_8-upload: # Uploading + wheel-py3_10-cuda13_0-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_10-cuda12_8-test + needs: wheel-py3_10-cuda13_0-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.10" - build_name: wheel-py3_10-cuda12_8 + build_name: wheel-py3_10-cuda13_0 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_10-cuda13_0-build: + wheel-py3_10-cuda13_2-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -748,8 +749,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.10" @@ -828,7 +829,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_10-cuda13_0 + name: wheel-py3_10-cuda13_2 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -844,10 +845,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_10-cuda13_0-test: # Testing + wheel-py3_10-cuda13_2-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_10-cuda13_0-build + - wheel-py3_10-cuda13_2-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -856,8 +857,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.10" @@ -928,7 +929,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_10-cuda13_0 + name: wheel-py3_10-cuda13_2 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -949,22 +950,22 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_10-cuda13_0-upload: # Uploading + wheel-py3_10-cuda13_2-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_10-cuda13_0-test + needs: wheel-py3_10-cuda13_2-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.10" - build_name: wheel-py3_10-cuda13_0 + build_name: wheel-py3_10-cuda13_2 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -1216,6 +1217,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.11" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -1665,7 +1667,7 @@ jobs: R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_11-cuda12_8-build: + wheel-py3_11-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -1675,8 +1677,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.11" @@ -1755,7 +1757,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_11-cuda12_8 + name: wheel-py3_11-cuda13_0 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -1771,10 +1773,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_11-cuda12_8-test: # Testing + wheel-py3_11-cuda13_0-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_11-cuda12_8-build + - wheel-py3_11-cuda13_0-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -1783,8 +1785,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.11" @@ -1855,7 +1857,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_11-cuda12_8 + name: wheel-py3_11-cuda13_0 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -1876,29 +1878,29 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_11-cuda12_8-upload: # Uploading + wheel-py3_11-cuda13_0-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_11-cuda12_8-test + needs: wheel-py3_11-cuda13_0-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.11" - build_name: wheel-py3_11-cuda12_8 + build_name: wheel-py3_11-cuda13_0 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_11-cuda13_0-build: + wheel-py3_11-cuda13_2-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -1908,8 +1910,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.11" @@ -1988,7 +1990,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_11-cuda13_0 + name: wheel-py3_11-cuda13_2 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -2004,10 +2006,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_11-cuda13_0-test: # Testing + wheel-py3_11-cuda13_2-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_11-cuda13_0-build + - wheel-py3_11-cuda13_2-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -2016,8 +2018,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.11" @@ -2088,7 +2090,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_11-cuda13_0 + name: wheel-py3_11-cuda13_2 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -2109,22 +2111,22 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_11-cuda13_0-upload: # Uploading + wheel-py3_11-cuda13_2-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_11-cuda13_0-test + needs: wheel-py3_11-cuda13_2-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.11" - build_name: wheel-py3_11-cuda13_0 + build_name: wheel-py3_11-cuda13_2 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -2376,6 +2378,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.12" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -2825,7 +2828,7 @@ jobs: R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_12-cuda12_8-build: + wheel-py3_12-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -2835,8 +2838,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.12" @@ -2915,7 +2918,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_12-cuda12_8 + name: wheel-py3_12-cuda13_0 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -2931,10 +2934,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_12-cuda12_8-test: # Testing + wheel-py3_12-cuda13_0-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_12-cuda12_8-build + - wheel-py3_12-cuda13_0-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -2943,8 +2946,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.12" @@ -3015,7 +3018,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_12-cuda12_8 + name: wheel-py3_12-cuda13_0 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -3036,29 +3039,29 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_12-cuda12_8-upload: # Uploading + wheel-py3_12-cuda13_0-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_12-cuda12_8-test + needs: wheel-py3_12-cuda13_0-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.12" - build_name: wheel-py3_12-cuda12_8 + build_name: wheel-py3_12-cuda13_0 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_12-cuda13_0-build: + wheel-py3_12-cuda13_2-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -3068,8 +3071,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.12" @@ -3148,7 +3151,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_12-cuda13_0 + name: wheel-py3_12-cuda13_2 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -3164,10 +3167,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_12-cuda13_0-test: # Testing + wheel-py3_12-cuda13_2-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_12-cuda13_0-build + - wheel-py3_12-cuda13_2-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -3176,8 +3179,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.12" @@ -3248,7 +3251,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_12-cuda13_0 + name: wheel-py3_12-cuda13_2 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -3269,22 +3272,22 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_12-cuda13_0-upload: # Uploading + wheel-py3_12-cuda13_2-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_12-cuda13_0-test + needs: wheel-py3_12-cuda13_2-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.12" - build_name: wheel-py3_12-cuda13_0 + build_name: wheel-py3_12-cuda13_2 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -3536,6 +3539,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -3985,7 +3989,7 @@ jobs: R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_13-cuda12_8-build: + wheel-py3_13-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -3995,8 +3999,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13" @@ -4075,7 +4079,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_13-cuda12_8 + name: wheel-py3_13-cuda13_0 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -4091,10 +4095,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_13-cuda12_8-test: # Testing + wheel-py3_13-cuda13_0-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_13-cuda12_8-build + - wheel-py3_13-cuda13_0-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -4103,8 +4107,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13" @@ -4175,7 +4179,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_13-cuda12_8 + name: wheel-py3_13-cuda13_0 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -4196,29 +4200,29 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_13-cuda12_8-upload: # Uploading + wheel-py3_13-cuda13_0-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_13-cuda12_8-test + needs: wheel-py3_13-cuda13_0-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.13" - build_name: wheel-py3_13-cuda12_8 + build_name: wheel-py3_13-cuda13_0 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_13-cuda13_0-build: + wheel-py3_13-cuda13_2-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -4228,8 +4232,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13" @@ -4308,7 +4312,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_13-cuda13_0 + name: wheel-py3_13-cuda13_2 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -4324,10 +4328,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_13-cuda13_0-test: # Testing + wheel-py3_13-cuda13_2-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_13-cuda13_0-build + - wheel-py3_13-cuda13_2-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -4336,8 +4340,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13" @@ -4408,7 +4412,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_13-cuda13_0 + name: wheel-py3_13-cuda13_2 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -4429,22 +4433,22 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_13-cuda13_0-upload: # Uploading + wheel-py3_13-cuda13_2-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_13-cuda13_0-test + needs: wheel-py3_13-cuda13_2-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.13" - build_name: wheel-py3_13-cuda13_0 + build_name: wheel-py3_13-cuda13_2 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -4696,6 +4700,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13t" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -5145,7 +5150,7 @@ jobs: R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_13t-cuda12_8-build: + wheel-py3_13t-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -5155,8 +5160,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13t" @@ -5235,7 +5240,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_13t-cuda12_8 + name: wheel-py3_13t-cuda13_0 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -5251,10 +5256,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_13t-cuda12_8-test: # Testing + wheel-py3_13t-cuda13_0-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_13t-cuda12_8-build + - wheel-py3_13t-cuda13_0-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -5263,8 +5268,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13t" @@ -5335,7 +5340,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_13t-cuda12_8 + name: wheel-py3_13t-cuda13_0 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -5356,29 +5361,29 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_13t-cuda12_8-upload: # Uploading + wheel-py3_13t-cuda13_0-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_13t-cuda12_8-test + needs: wheel-py3_13t-cuda13_0-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.13t" - build_name: wheel-py3_13t-cuda12_8 + build_name: wheel-py3_13t-cuda13_0 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_13t-cuda13_0-build: + wheel-py3_13t-cuda13_2-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -5388,8 +5393,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13t" @@ -5468,7 +5473,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_13t-cuda13_0 + name: wheel-py3_13t-cuda13_2 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -5484,10 +5489,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_13t-cuda13_0-test: # Testing + wheel-py3_13t-cuda13_2-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_13t-cuda13_0-build + - wheel-py3_13t-cuda13_2-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -5496,8 +5501,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.13t" @@ -5568,7 +5573,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_13t-cuda13_0 + name: wheel-py3_13t-cuda13_2 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -5589,22 +5594,22 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_13t-cuda13_0-upload: # Uploading + wheel-py3_13t-cuda13_2-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_13t-cuda13_0-test + needs: wheel-py3_13t-cuda13_2-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.13t" - build_name: wheel-py3_13t-cuda13_0 + build_name: wheel-py3_13t-cuda13_2 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -5856,6 +5861,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -6305,7 +6311,7 @@ jobs: R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_14-cuda12_8-build: + wheel-py3_14-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -6315,8 +6321,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14" @@ -6395,7 +6401,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_14-cuda12_8 + name: wheel-py3_14-cuda13_0 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -6411,10 +6417,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_14-cuda12_8-test: # Testing + wheel-py3_14-cuda13_0-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_14-cuda12_8-build + - wheel-py3_14-cuda13_0-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -6423,8 +6429,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14" @@ -6495,7 +6501,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_14-cuda12_8 + name: wheel-py3_14-cuda13_0 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -6516,29 +6522,29 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_14-cuda12_8-upload: # Uploading + wheel-py3_14-cuda13_0-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_14-cuda12_8-test + needs: wheel-py3_14-cuda13_0-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.14" - build_name: wheel-py3_14-cuda12_8 + build_name: wheel-py3_14-cuda13_0 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_14-cuda13_0-build: + wheel-py3_14-cuda13_2-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -6548,8 +6554,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14" @@ -6628,7 +6634,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_14-cuda13_0 + name: wheel-py3_14-cuda13_2 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -6644,10 +6650,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_14-cuda13_0-test: # Testing + wheel-py3_14-cuda13_2-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_14-cuda13_0-build + - wheel-py3_14-cuda13_2-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -6656,8 +6662,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14" @@ -6728,7 +6734,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_14-cuda13_0 + name: wheel-py3_14-cuda13_2 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -6749,22 +6755,22 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_14-cuda13_0-upload: # Uploading + wheel-py3_14-cuda13_2-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_14-cuda13_0-test + needs: wheel-py3_14-cuda13_2-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.14" - build_name: wheel-py3_14-cuda13_0 + build_name: wheel-py3_14-cuda13_2 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -7016,6 +7022,7 @@ jobs: GPU_ARCH_TYPE: cpu SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14t" + PYTORCH_EXTRA_INSTALL_REQUIREMENTS: cuda-toolkit[nvrtc,cudart,cupti,cufft,curand,cusolver,cusparse,cufile,nvjitlink,nvtx]==13.0.2; platform_system == 'Linux' | nvidia-cublas>=13.1.0.3,<=13.1.1.3; platform_system == 'Linux' | cuda-bindings>=13.0.3,<14; platform_system == 'Linux' | nvidia-cudnn-cu13==9.20.0.48; platform_system == 'Linux' | nvidia-cusparselt-cu13==0.8.1; platform_system == 'Linux' | nvidia-nccl-cu13==2.29.7; platform_system == 'Linux' | nvidia-nvshmem-cu13==3.4.5; platform_system == 'Linux' steps: # NOTE: These environment variables are put here so that they can be applied on every job equally # They are also here because setting them at a workflow level doesn't give us access to the @@ -7465,7 +7472,7 @@ jobs: R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_14t-cuda12_8-build: + wheel-py3_14t-cuda13_0-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -7475,8 +7482,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14t" @@ -7555,7 +7562,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_14t-cuda12_8 + name: wheel-py3_14t-cuda13_0 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -7571,10 +7578,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_14t-cuda12_8-test: # Testing + wheel-py3_14t-cuda13_0-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_14t-cuda12_8-build + - wheel-py3_14t-cuda13_0-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -7583,8 +7590,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14t" @@ -7655,7 +7662,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_14t-cuda12_8 + name: wheel-py3_14t-cuda13_0 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -7676,29 +7683,29 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_14t-cuda12_8-upload: # Uploading + wheel-py3_14t-cuda13_0-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_14t-cuda12_8-test + needs: wheel-py3_14t-cuda13_0-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.14t" - build_name: wheel-py3_14t-cuda12_8 + build_name: wheel-py3_14t-cuda13_0 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - wheel-py3_14t-cuda13_0-build: + wheel-py3_14t-cuda13_2-build: if: ${{ github.repository_owner == 'pytorch' }} needs: get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.12xlarge" @@ -7708,8 +7715,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14t" @@ -7788,7 +7795,7 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: wheel-py3_14t-cuda13_0 + name: wheel-py3_14t-cuda13_2 retention-days: 14 if-no-files-found: error path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" @@ -7804,10 +7811,10 @@ jobs: run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_14t-cuda13_0-test: # Testing + wheel-py3_14t-cuda13_2-test: # Testing if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_14t-cuda13_0-build + - wheel-py3_14t-cuda13_2-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}windows.g4dn.xlarge" timeout-minutes: 360 @@ -7816,8 +7823,8 @@ jobs: PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda SKIP_ALL_TESTS: 1 DESIRED_PYTHON: "3.14t" @@ -7888,7 +7895,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Build Artifacts with: - name: wheel-py3_14t-cuda13_0 + name: wheel-py3_14t-cuda13_2 path: "${{ env.PYTORCH_FINAL_PACKAGE_DIR }}" - name: Populate binary env shell: bash @@ -7909,22 +7916,22 @@ jobs: if: always() run: | .github\scripts\kill_active_ssh_sessions.ps1 - wheel-py3_14t-cuda13_0-upload: # Uploading + wheel-py3_14t-cuda13_2-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: wheel-py3_14t-cuda13_0-test + needs: wheel-py3_14t-cuda13_2-test with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: wheel # TODO: This is a legacy variable that we eventually want to get rid of in # favor of GPU_ARCH_VERSION - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda DESIRED_PYTHON: "3.14t" - build_name: wheel-py3_14t-cuda13_0 + build_name: wheel-py3_14t-cuda13_2 secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} @@ -8283,15 +8290,15 @@ jobs: R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - libtorch-cuda12_8-shared-with-deps-release-extract: + libtorch-cuda13_0-shared-with-deps-release-extract: if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_10-cuda12_8-build + - wheel-py3_10-cuda13_0-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" timeout-minutes: 60 env: - DESIRED_CUDA: cu128 + DESIRED_CUDA: cu130 LIBTORCH_VARIANT: shared-with-deps steps: - name: Checkout PyTorch @@ -8303,7 +8310,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Wheel Artifact with: - name: wheel-py3_10-cuda12_8 + name: wheel-py3_10-cuda13_0 path: "${{ runner.temp }}/wheel_artifact/" - name: Extract libtorch from wheel shell: bash @@ -8319,40 +8326,40 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: libtorch-cuda12_8-shared-with-deps-release + name: libtorch-cuda13_0-shared-with-deps-release retention-days: 14 if-no-files-found: error path: "${{ runner.temp }}/libtorch_output/" - libtorch-cuda12_8-shared-with-deps-release-upload: # Uploading + libtorch-cuda13_0-shared-with-deps-release-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: libtorch-cuda12_8-shared-with-deps-release-extract + needs: libtorch-cuda13_0-shared-with-deps-release-extract with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: libtorch - DESIRED_CUDA: cu128 - GPU_ARCH_VERSION: "12.8" + DESIRED_CUDA: cu130 + GPU_ARCH_VERSION: "13.0" GPU_ARCH_TYPE: cuda LIBTORCH_CONFIG: release LIBTORCH_VARIANT: shared-with-deps - build_name: libtorch-cuda12_8-shared-with-deps-release + build_name: libtorch-cuda13_0-shared-with-deps-release secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} R2_ACCESS_KEY_ID: ${{ secrets.R2_ACCESS_KEY_ID }} R2_SECRET_ACCESS_KEY: ${{ secrets.R2_SECRET_ACCESS_KEY }} uses: ./.github/workflows/_binary-upload.yml - libtorch-cuda13_0-shared-with-deps-release-extract: + libtorch-cuda13_2-shared-with-deps-release-extract: if: ${{ github.repository_owner == 'pytorch' }} needs: - - wheel-py3_10-cuda13_0-build + - wheel-py3_10-cuda13_2-build - get-label-type runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" timeout-minutes: 60 env: - DESIRED_CUDA: cu130 + DESIRED_CUDA: cu132 LIBTORCH_VARIANT: shared-with-deps steps: - name: Checkout PyTorch @@ -8364,7 +8371,7 @@ jobs: - uses: actions/download-artifact@v4.1.7 name: Download Wheel Artifact with: - name: wheel-py3_10-cuda13_0 + name: wheel-py3_10-cuda13_2 path: "${{ runner.temp }}/wheel_artifact/" - name: Extract libtorch from wheel shell: bash @@ -8380,25 +8387,25 @@ jobs: - uses: actions/upload-artifact@v4.4.0 if: always() with: - name: libtorch-cuda13_0-shared-with-deps-release + name: libtorch-cuda13_2-shared-with-deps-release retention-days: 14 if-no-files-found: error path: "${{ runner.temp }}/libtorch_output/" - libtorch-cuda13_0-shared-with-deps-release-upload: # Uploading + libtorch-cuda13_2-shared-with-deps-release-upload: # Uploading if: ${{ github.repository_owner == 'pytorch' }} permissions: id-token: write contents: read - needs: libtorch-cuda13_0-shared-with-deps-release-extract + needs: libtorch-cuda13_2-shared-with-deps-release-extract with: PYTORCH_ROOT: ${{ github.workspace }} PACKAGE_TYPE: libtorch - DESIRED_CUDA: cu130 - GPU_ARCH_VERSION: "13.0" + DESIRED_CUDA: cu132 + GPU_ARCH_VERSION: "13.2" GPU_ARCH_TYPE: cuda LIBTORCH_CONFIG: release LIBTORCH_VARIANT: shared-with-deps - build_name: libtorch-cuda13_0-shared-with-deps-release + build_name: libtorch-cuda13_2-shared-with-deps-release secrets: github-token: ${{ secrets.GITHUB_TOKEN }} R2_ACCOUNT_ID: ${{ secrets.R2_ACCOUNT_ID }} diff --git a/.github/workflows/h100-cutlass-backend.yml b/.github/workflows/h100-cutlass-backend.yml index e5406f7600133..3e02f3c80f57d 100644 --- a/.github/workflows/h100-cutlass-backend.yml +++ b/.github/workflows/h100-cutlass-backend.yml @@ -21,6 +21,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: diff --git a/.github/workflows/h100-distributed.yml b/.github/workflows/h100-distributed.yml index 0e5370a51c160..0a69cce8f868c 100644 --- a/.github/workflows/h100-distributed.yml +++ b/.github/workflows/h100-distributed.yml @@ -18,6 +18,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: diff --git a/.github/workflows/h100-symm-mem.yml b/.github/workflows/h100-symm-mem.yml index 09c362a546024..7ffe701656cb0 100644 --- a/.github/workflows/h100-symm-mem.yml +++ b/.github/workflows/h100-symm-mem.yml @@ -18,6 +18,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: diff --git a/.github/workflows/inductor-micro-benchmark-x86.yml b/.github/workflows/inductor-micro-benchmark-x86.yml index b0e9c869cdcf0..744ef7d7f75ce 100644 --- a/.github/workflows/inductor-micro-benchmark-x86.yml +++ b/.github/workflows/inductor-micro-benchmark-x86.yml @@ -16,6 +16,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: inductor-build: diff --git a/.github/workflows/inductor-micro-benchmark.yml b/.github/workflows/inductor-micro-benchmark.yml index 2473b5a515f77..8e2e6e04acbc7 100644 --- a/.github/workflows/inductor-micro-benchmark.yml +++ b/.github/workflows/inductor-micro-benchmark.yml @@ -16,6 +16,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-default-label-prefix: diff --git a/.github/workflows/inductor-nightly.yml b/.github/workflows/inductor-nightly.yml index 4258e8fdb0c84..24d0ebf9ef3f0 100644 --- a/.github/workflows/inductor-nightly.yml +++ b/.github/workflows/inductor-nightly.yml @@ -19,6 +19,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-default-label-prefix: diff --git a/.github/workflows/inductor-pallas.yml b/.github/workflows/inductor-pallas.yml index 8676434d0e580..72467be2fd9c6 100644 --- a/.github/workflows/inductor-pallas.yml +++ b/.github/workflows/inductor-pallas.yml @@ -16,6 +16,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: @@ -46,9 +47,6 @@ jobs: secrets: inherit linux-jammy-py3_12-inductor-pallas-gpu-test: - permissions: - id-token: write - contents: read name: pallas-gpu-py3.12-inductor uses: ./.github/workflows/_linux-test.yml needs: linux-jammy-py3_12-inductor-pallas-gpu-build @@ -73,9 +71,6 @@ jobs: secrets: inherit linux-jammy-py3_12-inductor-pallas-tpu-test: - permissions: - id-token: write - contents: read name: pallas-tpu-py3.12-inductor uses: ./.github/workflows/_linux-test.yml needs: linux-jammy-py3_12-inductor-pallas-tpu-build diff --git a/.github/workflows/inductor-perf-compare.yml b/.github/workflows/inductor-perf-compare.yml index 1ab538f96c1f0..f4b743058ce8a 100644 --- a/.github/workflows/inductor-perf-compare.yml +++ b/.github/workflows/inductor-perf-compare.yml @@ -13,6 +13,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-default-label-prefix: @@ -59,4 +60,5 @@ jobs: disable-monitor: false monitor-log-interval: 15 monitor-data-collect-interval: 4 + export-profiler-trace: "1" secrets: inherit diff --git a/.github/workflows/inductor-perf-test-b200.yml b/.github/workflows/inductor-perf-test-b200.yml index 96608d6d1d1d0..6eacb939c941a 100644 --- a/.github/workflows/inductor-perf-test-b200.yml +++ b/.github/workflows/inductor-perf-test-b200.yml @@ -64,6 +64,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: diff --git a/.github/workflows/inductor-perf-test-nightly-aarch64.yml b/.github/workflows/inductor-perf-test-nightly-aarch64.yml index f7b3517dccc06..d14244403998b 100644 --- a/.github/workflows/inductor-perf-test-nightly-aarch64.yml +++ b/.github/workflows/inductor-perf-test-nightly-aarch64.yml @@ -51,6 +51,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: diff --git a/.github/workflows/inductor-perf-test-nightly-h100.yml b/.github/workflows/inductor-perf-test-nightly-h100.yml index 4d161a00825b0..1075da051378d 100644 --- a/.github/workflows/inductor-perf-test-nightly-h100.yml +++ b/.github/workflows/inductor-perf-test-nightly-h100.yml @@ -51,6 +51,11 @@ on: required: false type: boolean default: false + deterministic_perf: + description: Run benchmarks with deterministic mode enabled? + required: false + type: boolean + default: false benchmark_configs: description: The list of configs used the benchmark required: false @@ -69,6 +74,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: @@ -131,7 +137,7 @@ jobs: if: github.event.schedule == '15 0 * * 1-6' with: build-environment: ${{ needs.build.outputs.build-environment }} - dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-cppwrapper-true-aotinductor-true-freezing_cudagraphs-true + dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-cppwrapper-true-aotinductor-true-freezing_cudagraphs-true-deterministic_perf-true docker-image: ${{ needs.build.outputs.docker-image }} test-matrix: ${{ needs.build.outputs.test-matrix }} timeout-minutes: 720 @@ -139,6 +145,8 @@ jobs: disable-monitor: false monitor-log-interval: 15 monitor-data-collect-interval: 4 + export-profiler-trace: "1" + enable-torch-trace: "1" secrets: inherit test-weekly: @@ -148,7 +156,7 @@ jobs: if: github.event.schedule == '0 7 * * 0' with: build-environment: ${{ needs.build.outputs.build-environment }} - dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-cppwrapper-true-aotinductor-true-freezing_cudagraphs-true-maxautotune-true-freeze_autotune_cudagraphs-true + dashboard-tag: training-true-inference-true-default-true-dynamic-true-cudagraphs-true-cppwrapper-true-aotinductor-true-freezing_cudagraphs-true-maxautotune-true-freeze_autotune_cudagraphs-true-deterministic_perf-true docker-image: ${{ needs.build.outputs.docker-image }} test-matrix: ${{ needs.build.outputs.test-matrix }} timeout-minutes: 1440 @@ -156,6 +164,8 @@ jobs: disable-monitor: false monitor-log-interval: 15 monitor-data-collect-interval: 4 + export-profiler-trace: "1" + enable-torch-trace: "1" secrets: inherit test: @@ -167,7 +177,7 @@ jobs: if: ${{ github.event_name == 'workflow_dispatch' || github.event_name == 'pull_request' }} with: build-environment: ${{ needs.build.outputs.build-environment }} - dashboard-tag: training-${{ inputs.training || 'true' }}-inference-${{ inputs.inference || 'true' }}-default-${{ inputs.default || 'true' }}-dynamic-${{ inputs.dynamic || 'true' }}-cudagraphs-${{ inputs.cudagraphs || 'true' }}-cppwrapper-${{ inputs.cppwrapper || 'false' }}-aotinductor-${{ inputs.aotinductor || 'false' }}-maxautotune-${{ inputs.maxautotune || 'false' }}-freezing_cudagraphs-${{ inputs.freezing_cudagraphs || 'false' }} + dashboard-tag: training-${{ inputs.training || 'true' }}-inference-${{ inputs.inference || 'true' }}-default-${{ inputs.default || 'true' }}-dynamic-${{ inputs.dynamic || 'true' }}-cudagraphs-${{ inputs.cudagraphs || 'true' }}-cppwrapper-${{ inputs.cppwrapper || 'false' }}-aotinductor-${{ inputs.aotinductor || 'false' }}-maxautotune-${{ inputs.maxautotune || 'false' }}-freezing_cudagraphs-${{ inputs.freezing_cudagraphs || 'false' }}-deterministic_perf-${{ inputs.deterministic_perf || 'false' }} docker-image: ${{ needs.build.outputs.docker-image }} test-matrix: ${{ needs.build.outputs.test-matrix }} timeout-minutes: 720 @@ -175,4 +185,6 @@ jobs: disable-monitor: false monitor-log-interval: 15 monitor-data-collect-interval: 4 + export-profiler-trace: "1" + enable-torch-trace: "1" secrets: inherit diff --git a/.github/workflows/inductor-perf-test-nightly-rocm-mi300.yml b/.github/workflows/inductor-perf-test-nightly-rocm-mi300.yml index c556c6b455783..9c92f11ad170a 100644 --- a/.github/workflows/inductor-perf-test-nightly-rocm-mi300.yml +++ b/.github/workflows/inductor-perf-test-nightly-rocm-mi300.yml @@ -79,10 +79,10 @@ jobs: linux-jammy-rocm-py3_10-inductor-benchmark-build: if: github.repository_owner == 'pytorch' - name: rocm-py3_10-inductor-benchmark-build + name: linux-jammy-rocm-py3.10-mi300 uses: ./.github/workflows/_linux-build.yml with: - build-environment: linux-jammy-rocm-py3_10 + build-environment: linux-jammy-rocm-py3.10-mi300 docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3-benchmarks test-matrix: | { include: [ @@ -114,7 +114,7 @@ jobs: permissions: id-token: write contents: read - name: rocm-py3_10-inductor-benchmark-test + name: linux-jammy-rocm-py3.10-mi300 uses: ./.github/workflows/_rocm-test.yml needs: linux-jammy-rocm-py3_10-inductor-benchmark-build with: diff --git a/.github/workflows/inductor-perf-test-nightly-rocm-mi355.yml b/.github/workflows/inductor-perf-test-nightly-rocm-mi355.yml index e6fd83193202c..f8b1ebd7dad9e 100644 --- a/.github/workflows/inductor-perf-test-nightly-rocm-mi355.yml +++ b/.github/workflows/inductor-perf-test-nightly-rocm-mi355.yml @@ -79,10 +79,10 @@ jobs: linux-jammy-rocm-py3_10-inductor-benchmark-build: if: github.repository_owner == 'pytorch' - name: rocm-py3_10-inductor-benchmark-build + name: linux-jammy-rocm-py3.10-mi355 uses: ./.github/workflows/_linux-build.yml with: - build-environment: linux-jammy-rocm-py3_10 + build-environment: linux-jammy-rocm-py3.10-mi355 docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3-benchmarks test-matrix: | { include: [ @@ -114,7 +114,7 @@ jobs: permissions: id-token: write contents: read - name: rocm-py3_10-inductor-benchmark-test + name: linux-jammy-rocm-py3.10-mi355 uses: ./.github/workflows/_rocm-test.yml needs: linux-jammy-rocm-py3_10-inductor-benchmark-build with: diff --git a/.github/workflows/inductor-perf-test-nightly-x86-zen.yml b/.github/workflows/inductor-perf-test-nightly-x86-zen.yml index eee51b7ff8889..a88d3d63ef50c 100644 --- a/.github/workflows/inductor-perf-test-nightly-x86-zen.yml +++ b/.github/workflows/inductor-perf-test-nightly-x86-zen.yml @@ -61,6 +61,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: diff --git a/.github/workflows/inductor-perf-test-nightly-x86.yml b/.github/workflows/inductor-perf-test-nightly-x86.yml index 87875831e2a0b..6c312586f27b5 100644 --- a/.github/workflows/inductor-perf-test-nightly-x86.yml +++ b/.github/workflows/inductor-perf-test-nightly-x86.yml @@ -61,6 +61,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: diff --git a/.github/workflows/inductor-perf-test-nightly.yml b/.github/workflows/inductor-perf-test-nightly.yml index 0d470a3bcd4cb..2e7c3d0e43f60 100644 --- a/.github/workflows/inductor-perf-test-nightly.yml +++ b/.github/workflows/inductor-perf-test-nightly.yml @@ -64,6 +64,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: diff --git a/.github/workflows/inductor-periodic.yml b/.github/workflows/inductor-periodic.yml index cfa844d3b7d96..a16a4fd80a63c 100644 --- a/.github/workflows/inductor-periodic.yml +++ b/.github/workflows/inductor-periodic.yml @@ -18,6 +18,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-default-label-prefix: @@ -59,6 +60,9 @@ jobs: { config: "dynamic_aot_eager_timm", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" }, { config: "dynamic_aot_eager_timm", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" }, { config: "dynamic_inductor_huggingface", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" }, + { config: "inductor_huggingface_unbacked_parity", shard: 1, num_shards: 1, runner: "linux.g5.4xlarge.nvidia.gpu" }, + { config: "inductor_huggingface_unbacked_parity", shard: 1, num_shards: 1, runner: "linux.aws.a100" }, + { config: "inductor_huggingface_unbacked_parity", shard: 1, num_shards: 1, runner: "linux.aws.h100" }, { config: "dynamic_inductor_timm", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" }, { config: "dynamic_inductor_timm", shard: 2, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" }, { config: "dynamic_inductor_torchbench", shard: 1, num_shards: 2, runner: "linux.g5.4xlarge.nvidia.gpu" }, @@ -87,7 +91,7 @@ jobs: name: rocm-periodic-dynamo-benchmarks-build uses: ./.github/workflows/_linux-build.yml with: - build-environment: linux-jammy-rocm-py3_10 + build-environment: linux-jammy-rocm-py3.10-mi355 docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3-benchmarks test-matrix: | { include: [ @@ -120,9 +124,6 @@ jobs: secrets: inherit rocm-periodic-dynamo-benchmarks-test: - permissions: - id-token: write - contents: read name: rocm-periodic-dynamo-benchmarks-test uses: ./.github/workflows/_rocm-test.yml needs: rocm-periodic-dynamo-benchmarks-build diff --git a/.github/workflows/inductor-rocm-mi200.yml b/.github/workflows/inductor-rocm-mi200.yml index c67f41daf5ad3..44fda5a4ed4e7 100644 --- a/.github/workflows/inductor-rocm-mi200.yml +++ b/.github/workflows/inductor-rocm-mi200.yml @@ -17,6 +17,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: @@ -31,12 +32,12 @@ jobs: opt_out_experiments: lf linux-jammy-rocm-py3_10-inductor-build: - name: rocm-py3.10-inductor + name: linux-jammy-rocm-py3.10-mi200 uses: ./.github/workflows/_linux-build.yml needs: get-label-type with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-rocm-py3.10 + build-environment: linux-jammy-rocm-py3.10-mi200 docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3 test-matrix: | { include: [ @@ -46,10 +47,7 @@ jobs: secrets: inherit linux-jammy-rocm-py3_10-inductor-test: - permissions: - id-token: write - contents: read - name: rocm-py3.10-inductor + name: linux-jammy-rocm-py3.10-mi200 uses: ./.github/workflows/_rocm-test.yml needs: linux-jammy-rocm-py3_10-inductor-build with: diff --git a/.github/workflows/inductor-rocm-mi300.yml b/.github/workflows/inductor-rocm-mi300.yml index 633386aba487b..0b196b4a214b9 100644 --- a/.github/workflows/inductor-rocm-mi300.yml +++ b/.github/workflows/inductor-rocm-mi300.yml @@ -17,15 +17,13 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: target-determination: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/target_determination.yml - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -39,7 +37,7 @@ jobs: opt_out_experiments: lf linux-noble-rocm-py3_12-inductor-build: - name: rocm-py3.12-inductor-mi300 + name: linux-noble-rocm-py3.12-mi300 uses: ./.github/workflows/_linux-build.yml needs: get-label-type with: @@ -54,10 +52,7 @@ jobs: secrets: inherit linux-noble-rocm-py3_12-inductor-test: - permissions: - id-token: write - contents: read - name: rocm-py3.12-inductor-mi300 + name: linux-noble-rocm-py3.12-mi300 uses: ./.github/workflows/_rocm-test.yml needs: linux-noble-rocm-py3_12-inductor-build with: diff --git a/.github/workflows/inductor-rocm-mi355.yml b/.github/workflows/inductor-rocm-mi355.yml index 70ea41a6da698..bd75369adbd00 100644 --- a/.github/workflows/inductor-rocm-mi355.yml +++ b/.github/workflows/inductor-rocm-mi355.yml @@ -16,15 +16,13 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: target-determination: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/target_determination.yml - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -38,7 +36,7 @@ jobs: opt_out_experiments: lf linux-noble-rocm-py3_12-inductor-build: - name: rocm-py3.12-inductor-mi355 + name: linux-noble-rocm-py3.12-mi355 uses: ./.github/workflows/_linux-build.yml needs: get-label-type with: @@ -53,10 +51,7 @@ jobs: secrets: inherit linux-noble-rocm-py3_12-inductor-test: - permissions: - id-token: write - contents: read - name: rocm-py3.12-inductor-mi355 + name: linux-noble-rocm-py3.12-mi355 uses: ./.github/workflows/_rocm-test.yml needs: linux-noble-rocm-py3_12-inductor-build with: diff --git a/.github/workflows/inductor-unittest.yml b/.github/workflows/inductor-unittest.yml index 82fb485f24b2b..cb5419437d703 100644 --- a/.github/workflows/inductor-unittest.yml +++ b/.github/workflows/inductor-unittest.yml @@ -18,6 +18,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: @@ -166,8 +167,8 @@ jobs: python-version: ['3.11', '3.12', '3.13'] with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-py${{ matrix.python-version }}-clang15 - docker-image-name: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang15 + build-environment: linux-jammy-py${{ matrix.python-version }}-clang18 + docker-image-name: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang18 test-matrix: | { include: [ { config: "inductor_core", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" }, @@ -183,8 +184,8 @@ jobs: matrix: python-version: ['3.11', '3.12', '3.13'] with: - build-environment: linux-jammy-py${{ matrix.python-version }}-clang15 - docker-image: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang15 + build-environment: linux-jammy-py${{ matrix.python-version }}-clang18 + docker-image: ci-image:pytorch-linux-jammy-py${{ matrix.python-version }}-clang18 test-matrix: | { include: [ { config: "inductor_core", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" }, diff --git a/.github/workflows/inductor.yml b/.github/workflows/inductor.yml index 2532279500849..d4a44968ffd93 100644 --- a/.github/workflows/inductor.yml +++ b/.github/workflows/inductor.yml @@ -25,6 +25,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: unit-test: @@ -73,6 +74,7 @@ jobs: build-environment: ${{ needs.inductor-build.outputs.build-environment }} docker-image: ${{ needs.inductor-build.outputs.docker-image }} test-matrix: ${{ needs.inductor-build.outputs.test-matrix }} + enable-torch-trace: "1" secrets: inherit inductor-cpu-build: diff --git a/.github/workflows/lint-osdc.yml b/.github/workflows/lint-osdc.yml deleted file mode 100644 index 1363bddf48dfe..0000000000000 --- a/.github/workflows/lint-osdc.yml +++ /dev/null @@ -1,373 +0,0 @@ -name: Lint OSDC (unstable) -# Unstable version of the lint workflow running on k8s ARC runners. -# Runs in parallel with the existing lint workflow during rollout. - -on: - pull_request: - branches-ignore: - - nightly - push: - branches: - - main - - release/* - - landchecks/* - tags: - - ciflow/pull/* - - ciflow/trunk/* - workflow_dispatch: - -permissions: read-all -# The names of steps that actually test the code should be suffixed with `(nonretryable)`. -# When any other step fails, it's job will be retried once by retryBot. -jobs: - get-label-type: - if: github.repository_owner == 'pytorch' - name: get-label-type - uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main - with: - triggering_actor: ${{ github.triggering_actor }} - issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }} - curr_branch: ${{ github.head_ref || github.ref_name }} - - get-changed-files: - if: github.repository_owner == 'pytorch' - name: Get changed files - uses: ./.github/workflows/_get-changed-files.yml - with: - all_files: ${{ contains(github.event.pull_request.labels.*.name, 'lint-all-files') || contains(github.event.pull_request.labels.*.name, 'Reverted') || github.event_name == 'push' }} - - lintrunner-clang: - uses: ./.github/workflows/_lint.yml - # Needed to prevent deduping on HUD - name: lintrunner-clang-${{ needs.get-changed-files.outputs.changed-files == '*' && 'all' || 'partial' }} - needs: [get-label-type, get-changed-files] - # Only run if there are changed files relevant to clangtidy / clangformat - if: | - github.repository_owner == 'pytorch' && ( - needs.get-changed-files.outputs.changed-files == '*' || - contains(needs.get-changed-files.outputs.changed-files, '.h') || - contains(needs.get-changed-files.outputs.changed-files, '.cpp') || - contains(needs.get-changed-files.outputs.changed-files, '.cc') || - contains(needs.get-changed-files.outputs.changed-files, '.cxx') || - contains(needs.get-changed-files.outputs.changed-files, '.hpp') || - contains(needs.get-changed-files.outputs.changed-files, '.hxx') || - contains(needs.get-changed-files.outputs.changed-files, '.cu') || - contains(needs.get-changed-files.outputs.changed-files, '.cuh') || - contains(needs.get-changed-files.outputs.changed-files, '.mm') || - contains(needs.get-changed-files.outputs.changed-files, '.metal') - ) - with: - runner: mt-l-x86iamx-8-16 - docker-image: ghcr.io/pytorch/test-infra:cuda-x86_64-67eb930 - script: | - CHANGED_FILES="${{ needs.get-changed-files.outputs.changed-files }}" - if [ "$CHANGED_FILES" = "*" ]; then - export ADDITIONAL_LINTRUNNER_ARGS="--take CLANGTIDY,CLANGFORMAT --all-files" - else - export ADDITIONAL_LINTRUNNER_ARGS="--take CLANGTIDY,CLANGFORMAT $CHANGED_FILES" - fi - export CLANG=1 - .github/scripts/lintrunner.sh - - # NOTE: mypy needs its own job because it depends on --all-files, without assessing all files it sometimes - # fails to find types when it should - # NOTE: We should be able to disable this and consolidate with Pyrefly - lintrunner-pyrefly: - uses: ./.github/workflows/_lint.yml - name: lintrunner-pyrefly-${{ needs.get-changed-files.outputs.changed-files == '*' && 'all' || 'partial' }} - needs: [get-label-type, get-changed-files] - # Only run if there are changed files relevant to pyrefly - if: | - github.repository_owner == 'pytorch' && ( - needs.get-changed-files.outputs.changed-files == '*' || - contains(needs.get-changed-files.outputs.changed-files, '.py') || - contains(needs.get-changed-files.outputs.changed-files, '.pyi') - ) - with: - runner: mt-l-x86iamx-8-16 - docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-67eb930 - script: | - CHANGED_FILES="${{ needs.get-changed-files.outputs.changed-files }}" - echo "Running pyrefly" - ADDITIONAL_LINTRUNNER_ARGS="--take PYREFLY --all-files" .github/scripts/lintrunner.sh - - lintrunner-noclang: - uses: ./.github/workflows/_lint.yml - name: lintrunner-noclang-${{ needs.get-changed-files.outputs.changed-files == '*' && 'all' || 'partial' }} - needs: [get-label-type, get-changed-files] - with: - runner: mt-l-x86iamx-8-16 - docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-67eb930 - script: | - CHANGED_FILES="${{ needs.get-changed-files.outputs.changed-files }}" - echo "Running all other linters" - if [ "$CHANGED_FILES" = '*' ]; then - ADDITIONAL_LINTRUNNER_ARGS="--skip CLANGTIDY,CLANGTIDY_EXECUTORCH_COMPATIBILITY,CLANGFORMAT,PYREFLY --all-files" .github/scripts/lintrunner.sh - else - ADDITIONAL_LINTRUNNER_ARGS="--skip CLANGTIDY,CLANGTIDY_EXECUTORCH_COMPATIBILITY,CLANGFORMAT,PYREFLY ${CHANGED_FILES}" .github/scripts/lintrunner.sh - fi - - quick-checks: - if: github.repository_owner == 'pytorch' - needs: get-label-type - uses: ./.github/workflows/_lint.yml - with: - runner: mt-l-x86iamx-8-16 - docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-67eb930 - script: | - # Ensure no non-breaking spaces - # NB: We use 'printf' below rather than '\u000a' since bash pre-4.2 - # does not support the '\u000a' syntax (which is relevant for local linters) - (! git --no-pager grep -In "$(printf '\xC2\xA0')" -- . || (echo "The above lines have non-breaking spaces (U+00A0); please convert them to spaces (U+0020)"; false)) - - # Ensure cross-OS compatible file names - (! git ls-files | grep -E '([<>:"|?*]|[ .]$)' || (echo "The above file names are not valid across all operating systems. Please ensure they don't contain the characters '<>:""|?*' and don't end with a white space or a '.' "; false)) - - # Ensure no versionless Python shebangs - (! git --no-pager grep -In '#!.*python$' -- . || (echo "The above lines have versionless Python shebangs; please specify either python2 or python3"; false)) - - # Ensure ciflow tags mentioned in config - python3 .github/scripts/collect_ciflow_labels.py --validate-tags - - # C++ docs check - pushd docs/cpp/source - ./check-doxygen.sh - popd - - # CUDA kernel launch check - set -eux - python3 torch/testing/_internal/check_kernel_launches.py |& tee cuda_kernel_launch_checks.txt - - pr-sanity-checks: - name: pr-sanity-checks - runs-on: linux.24_04.4x - # Only run this on pull requests. This check is simple enough to be done without a Docker image - if: ${{ github.event_name == 'pull_request' && !contains(github.event.pull_request.labels.*.name, 'skip-pr-sanity-checks') && github.repository_owner == 'pytorch' }} - steps: - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - with: - submodules: false - fetch-depth: -1 - - - name: PR size check (nonretryable) - env: - BASE: ${{ github.event.pull_request.base.sha }} - HEAD: ${{ github.event.pull_request.head.sha }} - run: | - bash .github/scripts/pr-sanity-check.sh - - workflow-checks: - if: github.repository_owner == 'pytorch' - needs: get-label-type - uses: ./.github/workflows/_lint.yml - with: - runner: mt-l-x86iamx-8-16 - docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-67eb930 - script: | - # Regenerate workflows - .github/scripts/generate_ci_workflows.py - - RC=0 - # Assert that regenerating the workflows didn't change them - if ! .github/scripts/report_git_status.sh .github/workflows; then - echo - echo 'As shown by the above diff, the committed .github/workflows' - echo 'are not up to date according to .github/templates.' - echo 'Please run this command, commit, and push again to your PR:' - echo - echo ' .github/scripts/generate_ci_workflows.py' - echo - echo 'If running that command does nothing, you may need to rebase' - echo 'onto a more recent commit from the PyTorch main branch.' - RC=1 - fi - - # Check that jobs will be cancelled - .github/scripts/ensure_actions_will_cancel.py - - exit $RC - - toc: - if: github.repository_owner == 'pytorch' - needs: get-label-type - uses: ./.github/workflows/_lint.yml - with: - runner: mt-l-x86iamx-8-16 - docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-67eb930 - script: | - # Regenerate ToCs and check that they didn't change - set -eu - - export PATH=~/.npm-global/bin:"$PATH" - for FILE in $(git grep -Il '' -- '**.md'); do - markdown-toc --bullets='-' -i "$FILE" - done - - if ! .github/scripts/report_git_status.sh .; then - echo - echo 'As shown by the above diff, the table of contents in one or' - echo 'more Markdown files is not up to date with the file contents.' - echo 'You can either apply that Git diff directly to correct the' - echo 'table of contents, or if you have npm installed, you can' - echo 'install the npm package markdown-toc and run the following' - # shellcheck disable=SC2016 - echo 'command (replacing $FILE with the filename for which you want' - echo 'to regenerate the table of contents):' - echo - # shellcheck disable=SC2016 - echo " markdown-toc --bullets='-' -i \"\$FILE\"" - false - fi - - test-tools: - name: Test tools - if: ${{ github.repository == 'pytorch/pytorch' }} - needs: get-label-type - uses: ./.github/workflows/_lint.yml - with: - runner: mt-l-x86iamx-8-16 - docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-67eb930 - script: | - # Test tools - PYTHONPATH=$(pwd) pytest tools/stats - PYTHONPATH=$(pwd) pytest tools/test -o "python_files=test*.py" - PYTHONPATH=$(pwd) pytest .github/scripts -o "python_files=test*.py" - - test_run_test: - name: Test `run_test.py` is usable without boto3 - if: ${{ github.repository == 'pytorch/pytorch' }} - runs-on: linux.24_04.4x - steps: - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - with: - submodules: false - fetch-depth: 1 - - name: Setup Python 3.10 - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: '3.10' - architecture: x64 - cache: pip - - name: Install dependencies - run: | - python3 -m pip install --upgrade pip - pip install pytest-rerunfailures==11.1.* pytest-flakefinder==1.1.* pytest-xdist==3.3.* expecttest==0.3.* fbscribelogger==0.1.* numpy==1.24.* - pip install torch --pre --index-url https://download.pytorch.org/whl/nightly/cpu/ - - name: Run run_test.py (nonretryable) - run: | - # Run test_vulkan, which is a fast noop on Linux - python3 test/run_test.py --include test_vulkan --verbose - - test_collect_env: - if: ${{ github.repository == 'pytorch/pytorch' }} - name: Test collect_env - runs-on: ${{ matrix.runner }} - strategy: - matrix: - include: - - test_type: with_torch - runner: linux.24_04.4x - - test_type: without_torch - runner: linux.24_04.4x - - test_type: older_python_version - runner: linux.24_04.4x - steps: - # [see note: pytorch repo ref] - # deep clone (fetch-depth 0) required, to allow us to use git log - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - with: - submodules: false - fetch-depth: 1 - - name: Get min python version - id: get-min-python-version - if: matrix.test_type == 'older_python_version' - run: | - set -eou pipefail - # Generate PyTorch version to use - echo "MIN_PYTHON_VERSION=$(python3 .github/scripts/get_ci_variable.py --min-python-version)" >> "${GITHUB_OUTPUT}" - - name: Setup Old Python version - if: matrix.test_type == 'older_python_version' - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: 3.8 - architecture: x64 - check-latest: false - cache: pip - cache-dependency-path: | - **/requirements-build.txt - **/requirements.txt - - name: Setup Min Python version - if: matrix.test_type != 'older_python_version' - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: ${{ steps.get-min-python-version.outputs.MIN_PYTHON_VERSION }} - architecture: x64 - check-latest: false - cache: pip - cache-dependency-path: | - **/requirements-build.txt - **/requirements.txt - - name: Install torch - if: matrix.test_type == 'with_torch' - run: | - pip install -r requirements.txt - # Doesn't really matter what torch version, we just need ANY torch installed - pip install 'torch==2.*' - - name: Run collect_env.py (nonretryable) - run: | - # All we need to see is that it passes - python3 torch/utils/collect_env.py - - link-check: - if: github.repository_owner == 'pytorch' - needs: get-label-type - name: Link checks - uses: ./.github/workflows/_link_check_osdc.yml - with: - runner: ${{ needs.get-label-type.outputs.label-type }} - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} - - doc-redirects-check: - name: doc-redirects-check - runs-on: linux.24_04.4x - if: github.event_name == 'pull_request' && github.repository_owner == 'pytorch' - steps: - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - with: - submodules: false - fetch-depth: 0 - - - name: Doc redirects check (nonretryable) - env: - BASE_REF: ${{ github.event.pull_request.base.ref }} - run: | - set -euo pipefail - - # Run the check with auto-fix to generate suggestions - python3 .github/scripts/check_doc_redirects.py \ - --base-ref "origin/${BASE_REF}" \ - --auto-fix - - # If redirects.py was modified, show the diff and fail - if ! git diff --quiet docs/source/redirects.py 2>/dev/null; then - echo "" - echo "📋 The following redirects were auto-generated:" - echo "" - git diff docs/source/redirects.py - echo "" - echo "Please add these changes to your PR by running locally:" - echo "" - echo " python3 .github/scripts/check_doc_redirects.py --base-ref origin/main --auto-fix" - echo " git add docs/source/redirects.py" - echo " git commit --amend --no-edit # or create a new commit" - echo "" - exit 1 - fi - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }} - cancel-in-progress: true diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index aab977b11469b..cd244274039d5 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -36,7 +36,7 @@ jobs: all_files: ${{ contains(github.event.pull_request.labels.*.name, 'lint-all-files') || contains(github.event.pull_request.labels.*.name, 'Reverted') || github.event_name == 'push' }} lintrunner-clang: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: ./.github/workflows/_lint.yml # Needed to prevent deduping on HUD name: lintrunner-clang-${{ needs.get-changed-files.outputs.changed-files == '*' && 'all' || 'partial' }} needs: [get-label-type, get-changed-files] @@ -56,14 +56,8 @@ jobs: contains(needs.get-changed-files.outputs.changed-files, '.metal') ) with: - timeout: 120 - runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" - docker-image: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-linter - # NB: A shallow checkout won't work here because calculate-docker-image requires a full checkout - # to run git rev-parse HEAD~:.ci/docker when a new image is needed - fetch-depth: 0 - submodules: true - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + runner: mt-l-x86iamx-8-16 + docker-image: ghcr.io/pytorch/test-infra:cuda-x86_64-810d48d script: | CHANGED_FILES="${{ needs.get-changed-files.outputs.changed-files }}" if [ "$CHANGED_FILES" = "*" ]; then @@ -78,7 +72,7 @@ jobs: # fails to find types when it should # NOTE: We should be able to disable this and consolidate with Pyrefly lintrunner-pyrefly: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: ./.github/workflows/_lint.yml name: lintrunner-pyrefly-${{ needs.get-changed-files.outputs.changed-files == '*' && 'all' || 'partial' }} needs: [get-label-type, get-changed-files] # Only run if there are changed files relevant to pyrefly @@ -89,32 +83,20 @@ jobs: contains(needs.get-changed-files.outputs.changed-files, '.pyi') ) with: - timeout: 120 - runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" - docker-image: ci-image:pytorch-linux-jammy-linter - # NB: A shallow checkout won't work here because calculate-docker-image requires a full checkout - # to run git rev-parse HEAD~:.ci/docker when a new image is needed - fetch-depth: 0 - submodules: true - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + runner: mt-l-x86iamx-8-16 + docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d script: | CHANGED_FILES="${{ needs.get-changed-files.outputs.changed-files }}" echo "Running pyrefly" ADDITIONAL_LINTRUNNER_ARGS="--take PYREFLY --all-files" .github/scripts/lintrunner.sh lintrunner-noclang: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: ./.github/workflows/_lint.yml name: lintrunner-noclang-${{ needs.get-changed-files.outputs.changed-files == '*' && 'all' || 'partial' }} needs: [get-label-type, get-changed-files] with: - timeout: 120 - runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" - docker-image: ci-image:pytorch-linux-jammy-linter - # NB: A shallow checkout won't work here because calculate-docker-image requires a full checkout - # to run git rev-parse HEAD~:.ci/docker when a new image is needed - fetch-depth: 0 - submodules: true - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + runner: mt-l-x86iamx-8-16 + docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d script: | CHANGED_FILES="${{ needs.get-changed-files.outputs.changed-files }}" echo "Running all other linters" @@ -125,14 +107,12 @@ jobs: fi quick-checks: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + if: github.repository_owner == 'pytorch' needs: get-label-type + uses: ./.github/workflows/_lint.yml with: - timeout: 120 - runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" - docker-image: ci-image:pytorch-linux-jammy-linter - fetch-depth: 0 - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + runner: mt-l-x86iamx-8-16 + docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d script: | # Ensure no non-breaking spaces # NB: We use 'printf' below rather than '\u000a' since bash pre-4.2 @@ -159,10 +139,9 @@ jobs: pr-sanity-checks: name: pr-sanity-checks - needs: get-label-type - runs-on: [self-hosted, "${{ needs.get-label-type.outputs.label-type }}linux.large"] + runs-on: linux.24_04.4x # Only run this on pull requests. This check is simple enough to be done without a Docker image - if: github.event_name == 'pull_request' && !contains(github.event.pull_request.labels.*.name, 'skip-pr-sanity-checks') + if: ${{ github.event_name == 'pull_request' && !contains(github.event.pull_request.labels.*.name, 'skip-pr-sanity-checks') && github.repository_owner == 'pytorch' }} steps: - name: Checkout PyTorch uses: pytorch/pytorch/.github/actions/checkout-pytorch@main @@ -177,17 +156,13 @@ jobs: run: | bash .github/scripts/pr-sanity-check.sh - workflow-checks: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + if: github.repository_owner == 'pytorch' needs: get-label-type + uses: ./.github/workflows/_lint.yml with: - timeout: 120 - runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" - docker-image: ci-image:pytorch-linux-jammy-linter - fetch-depth: -1 - submodules: true - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + runner: mt-l-x86iamx-8-16 + docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d script: | # Regenerate workflows .github/scripts/generate_ci_workflows.py @@ -213,14 +188,12 @@ jobs: exit $RC toc: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + if: github.repository_owner == 'pytorch' needs: get-label-type + uses: ./.github/workflows/_lint.yml with: - timeout: 120 - runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" - docker-image: ci-image:pytorch-linux-jammy-linter - fetch-depth: 0 - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + runner: mt-l-x86iamx-8-16 + docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d script: | # Regenerate ToCs and check that they didn't change set -eu @@ -249,14 +222,11 @@ jobs: test-tools: name: Test tools if: ${{ github.repository == 'pytorch/pytorch' }} - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main needs: get-label-type + uses: ./.github/workflows/_lint.yml with: - timeout: 120 - runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" - docker-image: ci-image:pytorch-linux-jammy-linter - fetch-depth: 0 - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + runner: mt-l-x86iamx-8-16 + docker-image: ghcr.io/pytorch/test-infra:cpu-x86_64-810d48d script: | # Test tools PYTHONPATH=$(pwd) pytest tools/stats @@ -351,18 +321,17 @@ jobs: python3 torch/utils/collect_env.py link-check: - name: Link checks + if: github.repository_owner == 'pytorch' needs: get-label-type + name: Link checks uses: ./.github/workflows/_link_check.yml with: runner: ${{ needs.get-label-type.outputs.label-type }} - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} - secrets: inherit + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} doc-redirects-check: name: doc-redirects-check - needs: get-label-type - runs-on: [self-hosted, "${{ needs.get-label-type.outputs.label-type }}linux.large"] + runs-on: linux.24_04.4x if: github.event_name == 'pull_request' && github.repository_owner == 'pytorch' steps: - name: Checkout PyTorch diff --git a/.github/workflows/linux-aarch64.yml b/.github/workflows/linux-aarch64.yml deleted file mode 100644 index 2c6e6b6dac39c..0000000000000 --- a/.github/workflows/linux-aarch64.yml +++ /dev/null @@ -1,62 +0,0 @@ -name: linux-aarch64 - -on: - push: - branches: - - main - - release/* - tags: - - ciflow/linux-aarch64/* - - ciflow/trunk/* - workflow_dispatch: - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} - cancel-in-progress: true - -jobs: - - get-label-type: - if: github.repository_owner == 'pytorch' - name: get-label-type - uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main - with: - triggering_actor: ${{ github.triggering_actor }} - issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }} - curr_branch: ${{ github.head_ref || github.ref_name }} - curr_ref_type: ${{ github.ref_type }} - - linux-jammy-aarch64-py3_10-build: - name: linux-jammy-aarch64-py3.10 - uses: ./.github/workflows/_linux-build.yml - needs: get-label-type - with: - runner_prefix: ${{ needs.get-label-type.outputs.label-type }} - build-environment: linux-jammy-aarch64-py3.10 - docker-image-name: ci-image:pytorch-linux-jammy-aarch64-py3.10-gcc13 - runner: linux.arm64.m7g.4xlarge - test-matrix: | - { include: [ - { config: "default", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" }, - { config: "default", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" }, - { config: "default", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" }, - { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" }, - { config: "default", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" }, - { config: "default", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" }, - { config: "default", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" }, - { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" }, - ]} - secrets: inherit - - linux-jammy-aarch64-py3_10-test: - name: linux-jammy-aarch64-py3.10 - uses: ./.github/workflows/_linux-test.yml - needs: linux-jammy-aarch64-py3_10-build - permissions: - id-token: write - contents: read - with: - build-environment: ${{ needs.linux-jammy-aarch64-py3_10-build.outputs.build-environment }} - docker-image: ${{ needs.linux-jammy-aarch64-py3_10-build.outputs.docker-image }} - test-matrix: ${{ needs.linux-jammy-aarch64-py3_10-build.outputs.test-matrix }} - secrets: inherit diff --git a/.github/workflows/llm_td_retrieval.yml b/.github/workflows/llm_td_retrieval.yml index 565a9b25df50f..23416b9b0627f 100644 --- a/.github/workflows/llm_td_retrieval.yml +++ b/.github/workflows/llm_td_retrieval.yml @@ -26,15 +26,8 @@ jobs: continue-on-error: true needs: get-label-type steps: - - name: Clone PyTorch - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - repository: pytorch/pytorch - fetch-depth: 0 - path: pytorch - - name: Setup Linux - uses: ./pytorch/.github/actions/setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main - name: Clone CodeLlama uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index c6282855aa88c..82f718c9b1077 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -18,6 +18,11 @@ concurrency: group: ${{ github.workflow }}--${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }} cancel-in-progress: true +permissions: + id-token: write + contents: read + actions: read + jobs: get-label-type: name: get-label-type @@ -45,7 +50,7 @@ jobs: with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build-environment: linux-jammy-py3.10-gcc11 - docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11 + docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18 secrets: inherit docs-push: diff --git a/.github/workflows/nitpicker.yml b/.github/workflows/nitpicker.yml index 40bd245ce913f..a112fb0dcec42 100644 --- a/.github/workflows/nitpicker.yml +++ b/.github/workflows/nitpicker.yml @@ -20,7 +20,7 @@ jobs: steps: - name: Checkout PyTorch uses: pytorch/pytorch/.github/actions/checkout-pytorch@main - - uses: ethanis/nitpicker@v1 + - uses: ethanis/nitpicker@c102a39683a80c7db9065f8eab7de8b58871f946 # v1 with: nitpicks: '.github/nitpicks.yml' token: "${{ secrets.GITHUB_TOKEN }}" diff --git a/.github/workflows/operator_benchmark.yml b/.github/workflows/operator_benchmark.yml index e682e1eb06c24..7d6cbe58b726b 100644 --- a/.github/workflows/operator_benchmark.yml +++ b/.github/workflows/operator_benchmark.yml @@ -28,6 +28,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: x86-opbenchmark-build: diff --git a/.github/workflows/operator_microbenchmark.yml b/.github/workflows/operator_microbenchmark.yml index 445cdcc4be04a..ded4e0a235a2f 100644 --- a/.github/workflows/operator_microbenchmark.yml +++ b/.github/workflows/operator_microbenchmark.yml @@ -17,6 +17,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-label-type: diff --git a/.github/workflows/periodic-rocm-mi200.yml b/.github/workflows/periodic-rocm-mi200.yml index 3107d6897e0d7..234c7933b0596 100644 --- a/.github/workflows/periodic-rocm-mi200.yml +++ b/.github/workflows/periodic-rocm-mi200.yml @@ -17,23 +17,18 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: llm-td: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/llm_td_retrieval.yml - permissions: - id-token: write - contents: read target-determination: name: before-test uses: ./.github/workflows/target_determination.yml needs: llm-td - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -46,12 +41,12 @@ jobs: curr_ref_type: ${{ github.ref_type }} linux-jammy-rocm-py3_10-build: - name: linux-jammy-rocm-py3.10 + name: linux-jammy-rocm-py3.10-mi200 uses: ./.github/workflows/_linux-build.yml needs: get-label-type with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-rocm-py3.10 + build-environment: linux-jammy-rocm-py3.10-mi200 docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3 test-matrix: | { include: [ @@ -62,10 +57,7 @@ jobs: secrets: inherit linux-jammy-rocm-py3_10-test: - permissions: - id-token: write - contents: read - name: linux-jammy-rocm-py3.10 + name: linux-jammy-rocm-py3.10-mi200 uses: ./.github/workflows/_rocm-test.yml needs: - linux-jammy-rocm-py3_10-build diff --git a/.github/workflows/periodic-rocm-mi300.yml b/.github/workflows/periodic-rocm-mi300.yml index 88da168926444..cfbe99812d8f6 100644 --- a/.github/workflows/periodic-rocm-mi300.yml +++ b/.github/workflows/periodic-rocm-mi300.yml @@ -14,24 +14,21 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }}-${{ github.event.schedule }} cancel-in-progress: true -permissions: read-all +permissions: + id-token: write + contents: read + actions: read jobs: llm-td: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/llm_td_retrieval.yml - permissions: - id-token: write - contents: read target-determination: name: before-test uses: ./.github/workflows/target_determination.yml needs: llm-td - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -60,9 +57,6 @@ jobs: secrets: inherit linux-noble-rocm-py3_12-test: - permissions: - id-token: write - contents: read name: linux-noble-rocm-py3.12-mi300 uses: ./.github/workflows/_rocm-test.yml needs: diff --git a/.github/workflows/periodic-rocm-mi355.yml b/.github/workflows/periodic-rocm-mi355.yml index 8c6063260861d..7fd9d36b13423 100644 --- a/.github/workflows/periodic-rocm-mi355.yml +++ b/.github/workflows/periodic-rocm-mi355.yml @@ -15,24 +15,21 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }}-${{ github.event.schedule }} cancel-in-progress: true -permissions: read-all +permissions: + id-token: write + contents: read + actions: read jobs: llm-td: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/llm_td_retrieval.yml - permissions: - id-token: write - contents: read target-determination: name: before-test uses: ./.github/workflows/target_determination.yml needs: llm-td - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -61,9 +58,6 @@ jobs: secrets: inherit linux-noble-rocm-py3_12-test: - permissions: - id-token: write - contents: read name: linux-noble-rocm-py3.12-mi355 uses: ./.github/workflows/_rocm-test.yml needs: diff --git a/.github/workflows/periodic.yml b/.github/workflows/periodic.yml index 09018f6cf1613..a36b7eae500b5 100644 --- a/.github/workflows/periodic.yml +++ b/.github/workflows/periodic.yml @@ -23,23 +23,18 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: llm-td: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/llm_td_retrieval.yml - permissions: - id-token: write - contents: read target-determination: name: before-test uses: ./.github/workflows/target_determination.yml needs: llm-td - permissions: - id-token: write - contents: read get-label-type: name: get-label-type diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index 4cbc0b5077360..2bcda9b8b3476 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -33,6 +33,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: # See job-filter.yml for rules on adding job filter conditions @@ -47,17 +48,11 @@ jobs: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/llm_td_retrieval.yml - permissions: - id-token: write - contents: read target-determination: name: before-test uses: ./.github/workflows/target_determination.yml needs: llm-td - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -67,6 +62,11 @@ jobs: triggering_actor: ${{ github.triggering_actor }} issue_owner: ${{ github.event.pull_request.user.login || github.event.issue.user.login }} curr_branch: ${{ github.head_ref || github.ref_name }} + check_experiments: arc,lf + + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-py3.10-gcc11 (build + test) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ linux-jammy-py3_10-gcc11-build: if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-gcc11 ') || contains(needs.job-filter.outputs.jobs, ' linux-docs ') }} @@ -78,24 +78,29 @@ jobs: with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build-environment: linux-jammy-py3.10-gcc11 - docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11 + docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18 test-matrix: | { include: [ - { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, - { config: "default", shard: 2, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, - { config: "default", shard: 3, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, - { config: "default", shard: 4, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, - { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, + { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, + { config: "default", shard: 2, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, + { config: "default", shard: 3, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, + { config: "default", shard: 4, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, + { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, { config: "docs_test", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, { config: "jit_legacy", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, { config: "backwards_compat", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" }, - { config: "distributed", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, - { config: "distributed", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, - { config: "distributed", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, + # Use m7i instead of m7a to match the existing (Intel Xeon) numerics. DTensor crossref tests like linalg.multi_dot have tight + # float32 tolerances sensitive to different FMA/reduction order across CPU vendors. + { config: "distributed", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge.amx" }, + { config: "distributed", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge.amx" }, + { config: "distributed", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge.amx" }, { config: "numpy_2_x", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" }, { config: "libtorch_agnostic_targetting", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: gcc11 secrets: inherit linux-jammy-py3_10-gcc11-test: @@ -106,22 +111,74 @@ jobs: - linux-jammy-py3_10-gcc11-build - target-determination - job-filter + - get-label-type with: build-environment: ${{ needs.linux-jammy-py3_10-gcc11-build.outputs.build-environment }} docker-image: ${{ needs.linux-jammy-py3_10-gcc11-build.outputs.docker-image }} test-matrix: ${{ needs.linux-jammy-py3_10-gcc11-build.outputs.test-matrix }} tests-to-include: ${{ github.event.inputs.tests-to-include || '' }} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: gcc11 secrets: inherit + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-aarch64-py3.10-gcc11 (build + test) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + linux-jammy-aarch64-py3_10-gcc13-build: + name: linux-jammy-aarch64-py3.10 + uses: ./.github/workflows/_linux-build.yml + needs: + - get-label-type + - job-filter + with: + runner_prefix: ${{ needs.get-label-type.outputs.label-type }} + build-environment: linux-jammy-aarch64-py3.10 + docker-image-name: ci-image:pytorch-linux-jammy-aarch64-py3.10-gcc13 + runner: linux.arm64.m8g.4xlarge + test-matrix: | + { include: [ + { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" }, + { config: "default", shard: 2, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" }, + { config: "default", shard: 3, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" }, + { config: "default", shard: 4, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" }, + { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" }, + { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m8g.4xlarge" }, + ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: gcc13 + secrets: inherit + + linux-jammy-aarch64-py3_10-gcc13-test: + name: linux-jammy-aarch64-py3.10 + uses: ./.github/workflows/_linux-test.yml + needs: + - linux-jammy-aarch64-py3_10-gcc13-build + - target-determination + - job-filter + - get-label-type + with: + build-environment: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.build-environment }} + docker-image: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.docker-image }} + test-matrix: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.test-matrix }} + tests-to-include: ${{ github.event.inputs.tests-to-include || '' }} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: gcc13 + secrets: inherit - linux-jammy-py3_14t-clang15-build: - name: linux-jammy-py3.14t-clang15 + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-py3.14t-clang18 (build + test) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + linux-jammy-py3_14t-clang18-build: + name: linux-jammy-py3.14t-clang18 uses: ./.github/workflows/_linux-build.yml needs: get-label-type with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-py3.14t-clang15 - docker-image-name: ci-image:pytorch-linux-jammy-py3.14t-clang15 + build-environment: linux-jammy-py3.14t-clang18 + docker-image-name: ci-image:pytorch-linux-jammy-py3.14t-clang18 test-matrix: | { include: [ { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, @@ -131,26 +188,39 @@ jobs: { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, { config: "crossref", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, { config: "crossref", shard: 2, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, - { config: "dynamo_wrapped", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, - { config: "dynamo_wrapped", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, - { config: "dynamo_wrapped", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, + # Free-threaded Python 3.14t has ~20-50% higher per-object memory overhead from biased reference counting and + # per-object locks, and test_nn (531 tests) under dynamo wrapping with compiled autograd consistently OOMs at 64GB. + { config: "dynamo_wrapped", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, + { config: "dynamo_wrapped", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, + { config: "dynamo_wrapped", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, { config: "einops", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.14t" + compiler: clang18 secrets: inherit - linux-jammy-py3_14t-clang15-test: - name: linux-jammy-py3.14t-clang15 + linux-jammy-py3_14t-clang18-test: + name: linux-jammy-py3.14t-clang18 uses: ./.github/workflows/_linux-test.yml needs: - - linux-jammy-py3_14t-clang15-build + - linux-jammy-py3_14t-clang18-build - target-determination + - get-label-type with: - build-environment: linux-jammy-py3.14t-clang15 - docker-image: ${{ needs.linux-jammy-py3_14t-clang15-build.outputs.docker-image }} - test-matrix: ${{ needs.linux-jammy-py3_14t-clang15-build.outputs.test-matrix }} + build-environment: linux-jammy-py3.14t-clang18 + docker-image: ${{ needs.linux-jammy-py3_14t-clang18-build.outputs.docker-image }} + test-matrix: ${{ needs.linux-jammy-py3_14t-clang18-build.outputs.test-matrix }} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.14t" + compiler: clang18 secrets: inherit + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-docs ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + linux-docs: if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-docs ') }} name: linux-docs @@ -158,11 +228,21 @@ jobs: needs: - linux-jammy-py3_10-gcc11-build - job-filter + - get-label-type with: build-environment: ${{ needs.linux-jammy-py3_10-gcc11-build.outputs.build-environment }} docker-image: ${{ needs.linux-jammy-py3_10-gcc11-build.outputs.docker-image }} + run-doxygen: true + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" + python-version: "3.10" + compiler: gcc11 secrets: inherit + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-py3.10-gcc11-no-ops (build only) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + linux-jammy-py3_10-gcc11-no-ops: if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-gcc11-no-ops ') }} name: linux-jammy-py3.10-gcc11-no-ops @@ -173,13 +253,20 @@ jobs: with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build-environment: linux-jammy-py3.10-gcc11-no-ops - docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11 + docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18 test-matrix: | { include: [ { config: "default", shard: 1, num_shards: 1 }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: gcc11 secrets: inherit + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-py3.10-clang18-asan (build + test) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + linux-jammy-py3_10-clang18-asan-build: if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang18-asan ') }} name: linux-jammy-py3.10-clang18-asan @@ -191,7 +278,7 @@ jobs: runner: linux.c7i.4xlarge runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build-environment: linux-jammy-py3.10-clang18-asan - docker-image-name: ci-image:pytorch-linux-jammy-py3-clang18-asan + docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18 test-matrix: | { include: [ { config: "default", shard: 1, num_shards: 7, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, @@ -203,7 +290,11 @@ jobs: { config: "default", shard: 7, num_shards: 7, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, ]} - sync-tag: asan-build + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: clang18 + # TODO (huydhn): Add this back once other workflow migrates + # sync-tag: asan-build secrets: inherit linux-jammy-py3_10-clang18-asan-test: @@ -214,58 +305,34 @@ jobs: - linux-jammy-py3_10-clang18-asan-build - target-determination - job-filter + - get-label-type with: build-environment: ${{ needs.linux-jammy-py3_10-clang18-asan-build.outputs.build-environment }} docker-image: ${{ needs.linux-jammy-py3_10-clang18-asan-build.outputs.docker-image }} test-matrix: ${{ needs.linux-jammy-py3_10-clang18-asan-build.outputs.test-matrix }} - sync-tag: asan-test tests-to-include: ${{ github.event.inputs.tests-to-include || '' }} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: clang18 + # TODO (huydhn): Add this back once other workflow migrates + # sync-tag: asan-test secrets: inherit - linux-jammy-py3_10-clang15-onnx-build: - if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang15-onnx ') }} - name: linux-jammy-py3.10-clang15-onnx - uses: ./.github/workflows/_linux-build.yml - needs: - - get-label-type - - job-filter - with: - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-py3.10-clang15-onnx - docker-image-name: ci-image:pytorch-linux-jammy-py3-clang15-onnx - test-matrix: | - { include: [ - { config: "default", shard: 1, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" }, - { config: "default", shard: 2, num_shards: 2, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" }, - ]} - secrets: inherit - - linux-jammy-py3_10-clang15-onnx-test: - if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang15-onnx ') }} - name: linux-jammy-py3.10-clang15-onnx - uses: ./.github/workflows/_linux-test.yml - needs: - - linux-jammy-py3_10-clang15-onnx-build - - target-determination - - job-filter - with: - build-environment: ${{ needs.linux-jammy-py3_10-clang15-onnx-build.outputs.build-environment }} - docker-image: ${{ needs.linux-jammy-py3_10-clang15-onnx-build.outputs.docker-image }} - test-matrix: ${{ needs.linux-jammy-py3_10-clang15-onnx-build.outputs.test-matrix }} - tests-to-include: ${{ github.event.inputs.tests-to-include || '' }} - secrets: inherit + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-py3.10-clang18 (build + test) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ - linux-jammy-py3_10-clang15-build: - if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang15 ') }} - name: linux-jammy-py3.10-clang15 + linux-jammy-py3_10-clang18-build: + if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang18 ') }} + name: linux-jammy-py3.10-clang18 uses: ./.github/workflows/_linux-build.yml needs: - get-label-type - job-filter with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-py3.10-clang15 - docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang15 + build-environment: linux-jammy-py3.10-clang18 + docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18 test-matrix: | { include: [ { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, @@ -280,35 +347,47 @@ jobs: { config: "dynamo_wrapped", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, { config: "einops", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, + { config: "onnx", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: clang18 secrets: inherit - linux-jammy-py3_10-clang15-test: - if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang15 ') }} - name: linux-jammy-py3.10-clang15 + linux-jammy-py3_10-clang18-test: + if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-clang18 ') }} + name: linux-jammy-py3.10-clang18 uses: ./.github/workflows/_linux-test.yml needs: - - linux-jammy-py3_10-clang15-build + - linux-jammy-py3_10-clang18-build - target-determination - job-filter + - get-label-type with: - build-environment: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.build-environment }} - docker-image: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.docker-image }} - test-matrix: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.test-matrix }} + build-environment: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.build-environment }} + docker-image: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.docker-image }} + test-matrix: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.test-matrix }} tests-to-include: ${{ github.event.inputs.tests-to-include || '' }} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: clang18 secrets: inherit - linux-jammy-py3_14-clang15-build: - if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.14-clang15 ') }} - name: linux-jammy-py3.14-clang15 + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-py3.14-clang18 (build + test) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + + linux-jammy-py3_14-clang18-build: + if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.14-clang18 ') }} + name: linux-jammy-py3.14-clang18 uses: ./.github/workflows/_linux-build.yml needs: - get-label-type - job-filter with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-py3.14-clang15 - docker-image-name: ci-image:pytorch-linux-jammy-py3.14-clang15 + build-environment: linux-jammy-py3.14-clang18 + docker-image-name: ci-image:pytorch-linux-jammy-py3.14-clang18 test-matrix: | { include: [ { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, @@ -324,39 +403,58 @@ jobs: { config: "einops", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, { config: "openreg", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.14" + compiler: clang18 secrets: inherit - linux-jammy-py3_14-clang15-test: - if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.14-clang15 ') }} - name: linux-jammy-py3.14-clang15 + linux-jammy-py3_14-clang18-test: + if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.14-clang18 ') }} + name: linux-jammy-py3.14-clang18 uses: ./.github/workflows/_linux-test.yml needs: - - linux-jammy-py3_14-clang15-build + - linux-jammy-py3_14-clang18-build - job-filter + - get-label-type with: - build-environment: ${{ needs.linux-jammy-py3_14-clang15-build.outputs.build-environment }} - docker-image: ${{ needs.linux-jammy-py3_14-clang15-build.outputs.docker-image }} - test-matrix: ${{ needs.linux-jammy-py3_14-clang15-build.outputs.test-matrix }} + build-environment: ${{ needs.linux-jammy-py3_14-clang18-build.outputs.build-environment }} + docker-image: ${{ needs.linux-jammy-py3_14-clang18-build.outputs.docker-image }} + test-matrix: ${{ needs.linux-jammy-py3_14-clang18-build.outputs.test-matrix }} tests-to-include: ${{ github.event.inputs.tests-to-include || '' }} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.14" + compiler: clang18 secrets: inherit - linux-jammy-cuda12_8-cudnn9-py3_10-clang15-build: - if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda12.8-cudnn9-py3.10-clang15 ') }} - name: linux-jammy-cuda12.8-cudnn9-py3.10-clang15 + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-cuda12.8-cudnn9-py3.10-clang18 (build only) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + + linux-jammy-cuda12_8-cudnn9-py3_10-clang18-build: + if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cuda12.8-cudnn9-py3.10-clang18 ') }} + name: linux-jammy-cuda12.8-cudnn9-py3.10-clang18 uses: ./.github/workflows/_linux-build.yml needs: - get-label-type - job-filter with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-cuda12.8-cudnn9-py3.10-clang15 - docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang15 + build-environment: linux-jammy-cuda12.8-cudnn9-py3.10-clang18 + docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3.10-clang18 test-matrix: | { include: [ { config: "default", shard: 1, num_shards: 1 }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: clang18 + cuda-version: "12.8" secrets: inherit + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-cpu-py3.10-gcc11-bazel-test (build + test) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + linux-jammy-cpu-py3_10-gcc11-bazel-test: if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-cpu-py3.10-gcc11-bazel-test ') }} name: linux-jammy-cpu-py3.10-gcc11-bazel-test @@ -365,7 +463,6 @@ jobs: - get-label-type - job-filter with: - runner: "${{ needs.get-label-type.outputs.label-type }}linux.large" build-environment: linux-jammy-cuda12.8-py3.10-gcc11-bazel-test docker-image-name: ci-image:pytorch-linux-jammy-cuda12.8-cudnn9-py3-gcc11 cuda-version: cpu @@ -373,8 +470,16 @@ jobs: { include: [ { config: "default", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" + python-version: "3.10" + compiler: gcc11 secrets: inherit + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-py3.10-gcc11-mobile-lightweight-dispatch (build only) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + linux-jammy-py3_10-gcc11-mobile-lightweight-dispatch-build: if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3.10-gcc11-mobile-lightweight-dispatch-build ') }} name: linux-jammy-py3.10-gcc11-mobile-lightweight-dispatch-build @@ -385,16 +490,23 @@ jobs: with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build-environment: linux-jammy-py3.10-gcc11-mobile-lightweight-dispatch-build - docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11 + docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18 build-generates-artifacts: false test-matrix: | { include: [ { config: "default", shard: 1, num_shards: 1 }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" + compiler: gcc11 secrets: inherit + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-rocm-py3.10 (build) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + linux-jammy-rocm-py3_10-build: - if: github.event_name == 'pull_request' || (needs.job-filter.outputs.jobs != '' && contains(needs.job-filter.outputs.jobs, ' linux-jammy-rocm-py3.10 ')) + if: ${{ github.event_name == 'pull_request' || (needs.job-filter.outputs.jobs != '' && contains(needs.job-filter.outputs.jobs, ' linux-jammy-rocm-py3.10 ')) }} # don't run build twice on main name: linux-jammy-rocm-py3.10 uses: ./.github/workflows/_linux-build.yml @@ -405,45 +517,21 @@ jobs: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build-environment: linux-jammy-rocm-py3.10 docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3 - sync-tag: rocm-build + # TODO (huydhn): Add this back once other workflow migrates + # sync-tag: rocm-build test-matrix: | { include: [ { config: "default", shard: 1, num_shards: 3, runner: "linux.rocm.gpu.2" }, { config: "default", shard: 2, num_shards: 3, runner: "linux.rocm.gpu.2" }, { config: "default", shard: 3, num_shards: 3, runner: "linux.rocm.gpu.2" }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" secrets: inherit - linux-jammy-cuda13_0-py3_10-gcc11-inductor-build: - if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' cuda13.0-py3.10-gcc11-sm75 ') }} - name: cuda13.0-py3.10-gcc11-sm75 - uses: ./.github/workflows/_linux-build.yml - needs: - - get-label-type - - job-filter - with: - runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm75 - docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3-gcc11-inductor-benchmarks - test-matrix: | - { include: [ - { config: "pr_time_benchmarks", shard: 1, num_shards: 1, runner: "linux.g4dn.metal.nvidia.gpu" }, - ]} - secrets: inherit - - linux-jammy-cuda13_0-py3_10-gcc11-inductor-test: - if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' cuda13.0-py3.10-gcc11-sm75 ') }} - name: cuda13.0-py3.10-gcc11-sm75 - uses: ./.github/workflows/_linux-test.yml - needs: - - linux-jammy-cuda13_0-py3_10-gcc11-inductor-build - - job-filter - with: - build-environment: linux-jammy-cuda13.0-py3.10-gcc11-sm75 - docker-image: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-inductor-build.outputs.docker-image }} - test-matrix: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-inductor-build.outputs.test-matrix }} - tests-to-include: ${{ github.event.inputs.tests-to-include || '' }} - secrets: inherit + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ linux-jammy-xpu-n-py3.10 (build) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ linux-jammy-xpu-n-py3_10-build: if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-xpu-n-py3.10 ') }} @@ -453,8 +541,6 @@ jobs: - get-label-type - job-filter with: - # This should sync with the build in xpu.yml but xpu uses a larger runner - # sync-tag: linux-xpu-n-build runner_prefix: ${{ needs.get-label-type.outputs.label-type }} runner: linux.c7i.4xlarge build-environment: linux-noble-xpu-n-py3.10 @@ -466,20 +552,29 @@ jobs: { config: "default", shard: 3, num_shards: 4, runner: "linux.idc.xpu" }, { config: "default", shard: 4, num_shards: 4, runner: "linux.idc.xpu" }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.10" secrets: inherit + # ╠══════════════════════════════════════════════════════════════════════╣ + # ║ dynamo-cpython (build + test) ║ + # ╠══════════════════════════════════════════════════════════════════════╣ + dynamo-cpython-build: name: dynamo-cpython-build uses: ./.github/workflows/_linux-build.yml needs: get-label-type with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-py3.13-clang15 - docker-image-name: ci-image:pytorch-linux-jammy-py3.13-clang15 + build-environment: linux-jammy-py3.13-clang18 + docker-image-name: ci-image:pytorch-linux-jammy-py3.13-clang18 test-matrix: | { include: [ { config: "dynamo_cpython", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" }, ]} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.13" + compiler: clang18 secrets: inherit dynamo-cpython-test: @@ -487,10 +582,10 @@ jobs: uses: ./.github/workflows/_linux-test.yml needs: [get-label-type, dynamo-cpython-build] with: - build-environment: linux-jammy-py3.13-clang15 - docker-image: ci-image:pytorch-linux-jammy-py3.13-clang15 - test-matrix: | - { include: [ - { config: "dynamo_cpython", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.c7i.2xlarge" }, - ]} + build-environment: linux-jammy-py3.13-clang18 + docker-image: ${{ needs.dynamo-cpython-build.outputs.docker-image }} + test-matrix: ${{ needs.dynamo-cpython-build.outputs.test-matrix }} + use-arc: ${{ needs.get-label-type.outputs.use-arc == 'true' }} + python-version: "3.13" + compiler: clang18 secrets: inherit diff --git a/.github/workflows/quantization-periodic.yml b/.github/workflows/quantization-periodic.yml index 8dd97ff9308db..3b7ac1fb62a67 100644 --- a/.github/workflows/quantization-periodic.yml +++ b/.github/workflows/quantization-periodic.yml @@ -16,6 +16,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-default-label-prefix: diff --git a/.github/workflows/rocm-mi200.yml b/.github/workflows/rocm-mi200.yml index 7ca06317eebe5..16c9e0097df5d 100644 --- a/.github/workflows/rocm-mi200.yml +++ b/.github/workflows/rocm-mi200.yml @@ -38,12 +38,12 @@ jobs: linux-jammy-rocm-py3_10-build: if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }} - name: linux-jammy-rocm-py3.10 + name: linux-jammy-rocm-py3.10-mi200 uses: ./.github/workflows/_linux-build.yml needs: get-label-type with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-rocm-py3.10 + build-environment: linux-jammy-rocm-py3.10-mi200 docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3 sync-tag: rocm-build test-matrix: | @@ -61,7 +61,7 @@ jobs: permissions: id-token: write contents: read - name: linux-jammy-rocm-py3.10 + name: linux-jammy-rocm-py3.10-mi200 uses: ./.github/workflows/_rocm-test.yml needs: - linux-jammy-rocm-py3_10-build diff --git a/.github/workflows/rocm-mi300.yml b/.github/workflows/rocm-mi300.yml index 9c2bae06f32bd..95eef38c17981 100644 --- a/.github/workflows/rocm-mi300.yml +++ b/.github/workflows/rocm-mi300.yml @@ -14,16 +14,16 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true -permissions: read-all +permissions: + id-token: write + contents: read + actions: read jobs: target-determination: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/target_determination.yml - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -56,9 +56,6 @@ jobs: secrets: inherit linux-noble-rocm-py3_12-test: - permissions: - id-token: write - contents: read name: linux-noble-rocm-py3.12-mi300 uses: ./.github/workflows/_rocm-test.yml needs: diff --git a/.github/workflows/rocm-mi355.yml b/.github/workflows/rocm-mi355.yml index 5a77695011f3c..6fac9dc83d503 100644 --- a/.github/workflows/rocm-mi355.yml +++ b/.github/workflows/rocm-mi355.yml @@ -12,16 +12,16 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true -permissions: read-all +permissions: + id-token: write + contents: read + actions: read jobs: target-determination: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/target_determination.yml - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -54,9 +54,6 @@ jobs: secrets: inherit linux-noble-rocm-py3_12-test: - permissions: - id-token: write - contents: read name: linux-noble-rocm-py3.12-mi355 uses: ./.github/workflows/_rocm-test.yml needs: diff --git a/.github/workflows/rocm-navi31.yml b/.github/workflows/rocm-navi31.yml index b1dd74c3ba8d4..7452762a43bfb 100644 --- a/.github/workflows/rocm-navi31.yml +++ b/.github/workflows/rocm-navi31.yml @@ -15,16 +15,16 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true -permissions: read-all +permissions: + id-token: write + contents: read + actions: read jobs: target-determination: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/target_determination.yml - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -38,12 +38,12 @@ jobs: linux-jammy-rocm-py3_10-build: if: ${{ (github.event_name != 'schedule' || github.repository == 'pytorch/pytorch') && github.repository_owner == 'pytorch' }} - name: linux-jammy-rocm-py3.10 + name: linux-jammy-rocm-py3.10-navi31 uses: ./.github/workflows/_linux-build.yml needs: get-label-type with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-rocm-py3.10 + build-environment: linux-jammy-rocm-py3.10-navi31 docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3 sync-tag: rocm-build test-matrix: | @@ -54,10 +54,7 @@ jobs: secrets: inherit linux-jammy-rocm-py3_10-test: - permissions: - id-token: write - contents: read - name: linux-jammy-rocm-py3_10 + name: linux-jammy-rocm-py3.10-navi31 uses: ./.github/workflows/_rocm-test.yml needs: - linux-jammy-rocm-py3_10-build diff --git a/.github/workflows/rocm-nightly.yml b/.github/workflows/rocm-nightly.yml index 3a575f9b5607d..e17a002a22de4 100644 --- a/.github/workflows/rocm-nightly.yml +++ b/.github/workflows/rocm-nightly.yml @@ -2,19 +2,30 @@ name: rocm-nightly on: push: - branches: - - nightly # Trigger when nightly branch is updated tags: - ciflow/rocm-nightly/* - workflow_dispatch: # Allow manual triggering + workflow_dispatch: + schedule: + - cron: 0 0 * * * # midnight UTC concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true -permissions: read-all +permissions: + id-token: write + contents: read + actions: read jobs: + target-determination: + if: github.repository_owner == 'pytorch' + name: before-test + uses: ./.github/workflows/target_determination.yml + permissions: + id-token: write + contents: read + get-label-type: name: get-label-type uses: pytorch/pytorch/.github/workflows/_runner-determinator.yml@main @@ -34,4 +45,28 @@ jobs: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build-environment: linux-noble-rocm-nightly-py3.12-gfx942 docker-image-name: ci-image:pytorch-linux-noble-rocm-nightly-py3 + test-matrix: | + { include: [ + { config: "default", shard: 1, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" }, + { config: "default", shard: 2, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" }, + { config: "default", shard: 3, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" }, + { config: "default", shard: 4, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" }, + { config: "default", shard: 5, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" }, + { config: "default", shard: 6, num_shards: 6, runner: "linux.rocm.gpu.gfx942.1" }, + ]} + secrets: inherit + + linux-noble-rocm-nightly-py3_12-test: + permissions: + id-token: write + contents: read + name: linux-noble-rocm-nightly-py3.12-gfx942 + uses: ./.github/workflows/_rocm-test.yml + needs: + - linux-noble-rocm-nightly-py3_12-build + - target-determination + with: + build-environment: ${{ needs.linux-noble-rocm-nightly-py3_12-build.outputs.build-environment }} + docker-image: ${{ needs.linux-noble-rocm-nightly-py3_12-build.outputs.docker-image }} + test-matrix: ${{ needs.linux-noble-rocm-nightly-py3_12-build.outputs.test-matrix }} secrets: inherit diff --git a/.github/workflows/runner-determinator-validator.yml b/.github/workflows/runner-determinator-validator.yml deleted file mode 100644 index 0710229a7b9ff..0000000000000 --- a/.github/workflows/runner-determinator-validator.yml +++ /dev/null @@ -1,41 +0,0 @@ -name: Validate Runner Determinator Script is in Sync - -on: - # Run on PRs when the runner-determinator script is updated to ensure it's copies are kept in sync - pull_request: - paths: - - .github/workflows/_runner-determinator.yml - - .github/workflows/runner-determinator-validator.yml - - .github/scripts/runner_determinator.py - workflow_dispatch: - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} - cancel-in-progress: true - -jobs: - check-runner-determinator: - if: github.repository_owner == 'pytorch' - runs-on: ubuntu-latest - - steps: - - name: Checkout repository - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - - - name: Run Hardcode runner-determinator script - id: hardcode-script - run: | - # Extract the script content from _runner-determinator.yml and skip the first 10 spaces of each line - script_content=$(awk '/cat < runner_determinator.py/{flag=1;next}/EOF$/{flag=0}flag{print substr($0, 11)}' .github/workflows/_runner-determinator.yml) - - # Write the extracted script content to runner_determinator.py - echo "$script_content" > runner_determinator_workflow.py - - - name: Compare runner-determinator script embedded in workflow with checked in script - run: | - # Compare the extracted runner_determinator script with the existing one - # If this check fails, then make sure the contents of .github/scripts/runner_determinator.py is in sync with the - # version embedded into .github/workflows/_runner-determinator.yml - diff runner_determinator_workflow.py .github/scripts/runner_determinator.py - # Fail the job if the scripts are not identical - continue-on-error: false diff --git a/.github/workflows/runner_determinator_script_sync.yaml b/.github/workflows/runner_determinator_script_sync.yaml deleted file mode 100644 index a5f52f6980f7c..0000000000000 --- a/.github/workflows/runner_determinator_script_sync.yaml +++ /dev/null @@ -1,43 +0,0 @@ -name: runner-determinator - -on: - workflow_dispatch: - pull_request: - branches: [main] - paths: - - .github/workflows/_runner-determinator.yaml - - .github/workflows/_runner_determinator_script_sync.yaml - - .github/workflows/scripts/runner_determinator.py - -jobs: - python-script-sync-check: - if: github.repository_owner == 'pytorch' - runs-on: ubuntu-latest - - steps: - - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - with: - sparse-checkout: | - .github - - - name: Extract the script from runner_determinator - run: | - # Runner determinator files - RUNNER_DETERMINATOR_WORKFLOW_FILE=.github/workflows/_runner-determinator.yml - RUNNER_DETERMINATOR_PYTHON_SCRIPT_FILE=.github/scripts/runner_determinator.py - - # Parse the job file, extract the script and run it, up to the final EOF, - # to generate the python file in the local folder - yq '.jobs.runner-determinator.steps[] | select(.id == "hardcode-script") | .run' \ - "${RUNNER_DETERMINATOR_WORKFLOW_FILE}" | sed '/^EOF$/q' | bash - - set +e - DIFF="$(diff "$(basename ${RUNNER_DETERMINATOR_PYTHON_SCRIPT_FILE})" ${RUNNER_DETERMINATOR_PYTHON_SCRIPT_FILE})" - IS_DIFF=$? - set -e - if [ $IS_DIFF -eq 0 ]; then - echo "Scripts are in sync! ^_^"; - else - echo -e "Scripts are *NOT* in sync:\n ${DIFF}"; - exit 1 - fi diff --git a/.github/workflows/s390x-periodic.yml b/.github/workflows/s390x-periodic.yml index bf94ee3bb00f1..f583a5b425a44 100644 --- a/.github/workflows/s390x-periodic.yml +++ b/.github/workflows/s390x-periodic.yml @@ -18,23 +18,18 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: llm-td: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/llm_td_retrieval.yml - permissions: - id-token: write - contents: read target-determination: name: before-test uses: ./.github/workflows/target_determination.yml needs: llm-td - permissions: - id-token: write - contents: read linux-manylinux-2_28-py3-cpu-s390x-build: if: github.repository_owner == 'pytorch' @@ -60,9 +55,6 @@ jobs: secrets: inherit linux-manylinux-2_28-py3-cpu-s390x-test: - permissions: - id-token: write - contents: read name: linux-manylinux-2_28-py3-cpu-s390x uses: ./.github/workflows/_linux-test.yml needs: diff --git a/.github/workflows/slow-rocm-mi200.yml b/.github/workflows/slow-rocm-mi200.yml index f1e3bcc00b2b5..dfc0f70a7c598 100644 --- a/.github/workflows/slow-rocm-mi200.yml +++ b/.github/workflows/slow-rocm-mi200.yml @@ -21,23 +21,18 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: llm-td: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/llm_td_retrieval.yml - permissions: - id-token: write - contents: read target-determination: name: before-test uses: ./.github/workflows/target_determination.yml needs: llm-td - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -50,12 +45,12 @@ jobs: curr_ref_type: ${{ github.ref_type }} linux-jammy-rocm-py3_10-build: - name: linux-jammy-rocm-py3.10 + name: linux-jammy-rocm-py3.10-mi200 uses: ./.github/workflows/_linux-build.yml needs: get-label-type with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-rocm-py3.10 + build-environment: linux-jammy-rocm-py3.10-mi200 docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3 sync-tag: rocm-build test-matrix: | @@ -66,10 +61,7 @@ jobs: secrets: inherit linux-jammy-rocm-py3_10-test: - permissions: - id-token: write - contents: read - name: linux-jammy-rocm-py3.10 + name: linux-jammy-rocm-py3.10-mi200 uses: ./.github/workflows/_rocm-test.yml needs: - linux-jammy-rocm-py3_10-build diff --git a/.github/workflows/slow.yml b/.github/workflows/slow.yml index 8da9c9bd219d5..ff1fda43d39ba 100644 --- a/.github/workflows/slow.yml +++ b/.github/workflows/slow.yml @@ -21,23 +21,18 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: llm-td: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/llm_td_retrieval.yml - permissions: - id-token: write - contents: read target-determination: name: before-test uses: ./.github/workflows/target_determination.yml needs: llm-td - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -107,14 +102,14 @@ jobs: test-matrix: ${{ needs.linux-jammy-cuda13_0-py3_10-gcc11-sm86-build.outputs.test-matrix }} secrets: inherit - linux-jammy-py3_10-clang15-build: - name: linux-jammy-py3.10-clang15 + linux-jammy-py3_10-clang18-build: + name: linux-jammy-py3.10-clang18 uses: ./.github/workflows/_linux-build.yml needs: get-label-type with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-py3.10-clang15 - docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang15 + build-environment: linux-jammy-py3.10-clang18 + docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18 test-matrix: | { include: [ { config: "slow", shard: 1, num_shards: 2, runner: "linux.2xlarge" }, @@ -122,16 +117,16 @@ jobs: ]} secrets: inherit - linux-jammy-py3_10-clang15-test: - name: linux-jammy-py3.10-clang15 + linux-jammy-py3_10-clang18-test: + name: linux-jammy-py3.10-clang18 uses: ./.github/workflows/_linux-test.yml needs: - - linux-jammy-py3_10-clang15-build + - linux-jammy-py3_10-clang18-build - target-determination with: - build-environment: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.build-environment }} - docker-image: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.docker-image }} - test-matrix: ${{ needs.linux-jammy-py3_10-clang15-build.outputs.test-matrix }} + build-environment: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.build-environment }} + docker-image: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.docker-image }} + test-matrix: ${{ needs.linux-jammy-py3_10-clang18-build.outputs.test-matrix }} secrets: inherit linux-jammy-py3_10-clang18-asan-build: @@ -142,7 +137,7 @@ jobs: runner: linux.c7i.4xlarge runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" build-environment: linux-jammy-py3.10-clang18-asan - docker-image-name: ci-image:pytorch-linux-jammy-py3-clang18-asan + docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18 test-matrix: | { include: [ { config: "slow", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.4xlarge" }, diff --git a/.github/workflows/target-determination-indexer.yml b/.github/workflows/target-determination-indexer.yml index 3438b1dd5ac57..41b6b54ae9b5b 100644 --- a/.github/workflows/target-determination-indexer.yml +++ b/.github/workflows/target-determination-indexer.yml @@ -31,7 +31,7 @@ jobs: path: pytorch - name: Setup Linux - uses: ./pytorch/.github/actions/setup-linux + uses: pytorch/pytorch/.github/actions/setup-linux@main - name: Login to ECR uses: ./pytorch/.github/actions/ecr-login diff --git a/.github/workflows/target_determination.yml b/.github/workflows/target_determination.yml index c712b11185a76..9594ce0bf8f63 100644 --- a/.github/workflows/target_determination.yml +++ b/.github/workflows/target_determination.yml @@ -22,18 +22,11 @@ jobs: runs-on: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" needs: get-label-type steps: - # [pytorch repo ref] - # Use a pytorch/pytorch reference instead of a reference to the local - # checkout because when we run this action we don't *have* a local - # checkout. In other cases you should prefer a local checkout. - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main + - name: Setup Linux + uses: pytorch/pytorch/.github/actions/setup-linux@main with: submodules: false - - name: Setup Linux - uses: ./.github/actions/setup-linux - - name: Get workflow job id id: get-job-id uses: ./.github/actions/get-workflow-job-id diff --git a/.github/workflows/test-b200.yml b/.github/workflows/test-b200.yml index 19dcb07c29844..e753dd27ca938 100644 --- a/.github/workflows/test-b200.yml +++ b/.github/workflows/test-b200.yml @@ -35,6 +35,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: diff --git a/.github/workflows/test-h100.yml b/.github/workflows/test-h100.yml index 4351b427b0b8a..c1e8c13797da7 100644 --- a/.github/workflows/test-h100.yml +++ b/.github/workflows/test-h100.yml @@ -22,6 +22,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: diff --git a/.github/workflows/torchbench.yml b/.github/workflows/torchbench.yml index 2b370f6083185..f14e66576fd96 100644 --- a/.github/workflows/torchbench.yml +++ b/.github/workflows/torchbench.yml @@ -13,6 +13,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: get-default-label-prefix: diff --git a/.github/workflows/torchtitan.yml b/.github/workflows/torchtitan.yml index 92bda75958e54..b12f980ec4202 100644 --- a/.github/workflows/torchtitan.yml +++ b/.github/workflows/torchtitan.yml @@ -16,6 +16,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: build: diff --git a/.github/workflows/trunk-rocm-sandbox.yml b/.github/workflows/trunk-rocm-sandbox.yml index ddeb4e5809aec..bef62c9fbe5bc 100644 --- a/.github/workflows/trunk-rocm-sandbox.yml +++ b/.github/workflows/trunk-rocm-sandbox.yml @@ -16,23 +16,18 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: llm-td: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/llm_td_retrieval.yml - permissions: - id-token: write - contents: read target-determination: name: before-test uses: ./.github/workflows/target_determination.yml needs: llm-td - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -70,9 +65,6 @@ jobs: secrets: inherit linux-jammy-rocm-py3_10-test: - permissions: - id-token: write - contents: read name: linux-jammy-rocm-py3.10 uses: ./.github/workflows/_rocm-test.yml needs: diff --git a/.github/workflows/trunk-tagging.yml b/.github/workflows/trunk-tagging.yml index d96f2de8366aa..27472f786794c 100644 --- a/.github/workflows/trunk-tagging.yml +++ b/.github/workflows/trunk-tagging.yml @@ -28,8 +28,8 @@ jobs: - name: Pre-checkout validation run: | # For workflow_dispatch, validate SHA format before checkout - if [ -n "${{ github.event.inputs.commit_sha }}" ]; then - COMMIT_SHA="${{ github.event.inputs.commit_sha }}" + if [ -n "${INPUT_COMMIT_SHA}" ]; then + COMMIT_SHA="${INPUT_COMMIT_SHA}" # Verify it's a well-formed SHA (40 hex characters) if ! echo "${COMMIT_SHA}" | grep -qE '^[a-f0-9]{40}$'; then @@ -42,6 +42,8 @@ jobs: echo "✅ Using current commit SHA - no pre-checkout validation needed" fi + env: + INPUT_COMMIT_SHA: ${{ github.event.inputs.commit_sha }} - name: Checkout repository uses: actions/checkout@v4 with: @@ -53,8 +55,8 @@ jobs: - name: Set commit SHA id: commit run: | - if [ -n "${{ github.event.inputs.commit_sha }}" ]; then - COMMIT_SHA="${{ github.event.inputs.commit_sha }}" + if [ -n "${INPUT_COMMIT_SHA}" ]; then + COMMIT_SHA="${INPUT_COMMIT_SHA}" else COMMIT_SHA="${{ github.sha }}" fi @@ -63,6 +65,8 @@ jobs: echo "tag_name=trunk/${COMMIT_SHA}" } >> "${GITHUB_OUTPUT}" + env: + INPUT_COMMIT_SHA: ${{ github.event.inputs.commit_sha }} - name: Validate commit SHA run: | COMMIT_SHA="${{ steps.commit.outputs.sha }}" @@ -74,7 +78,7 @@ jobs: fi # For workflow_dispatch, verify the commit exists on main branch - if [ -n "${{ github.event.inputs.commit_sha }}" ]; then + if [ -n "${INPUT_COMMIT_SHA}" ]; then echo "Manual dispatch detected - validating commit is on main branch..." # Get all commits reachable from main branch @@ -89,6 +93,8 @@ jobs: echo "✅ Commit ${COMMIT_SHA} is valid (automatic push trigger)" fi + env: + INPUT_COMMIT_SHA: ${{ github.event.inputs.commit_sha }} - name: Create and push tag(s) with retry id: check_tag env: @@ -309,7 +315,10 @@ jobs: echo " Name: ${{ steps.commit.outputs.tag_name }}" echo " Commit: ${{ steps.commit.outputs.sha }}" echo " Trigger: ${{ github.event_name }}" - if [ -n "${{ github.event.inputs.commit_sha }}" ]; then - echo " Manual commit: ${{ github.event.inputs.commit_sha }}" + if [ -n "${INPUT_COMMIT_SHA}" ]; then + echo " Manual commit: ${INPUT_COMMIT_SHA}" fi fi + + env: + INPUT_COMMIT_SHA: ${{ github.event.inputs.commit_sha }} \ No newline at end of file diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml index e3bf1d63613dc..5460e803bd0db 100644 --- a/.github/workflows/trunk.yml +++ b/.github/workflows/trunk.yml @@ -30,6 +30,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: # See job-filter.yml for rules on adding job filter conditions @@ -44,17 +45,11 @@ jobs: if: github.repository_owner == 'pytorch' name: before-test uses: ./.github/workflows/llm_td_retrieval.yml - permissions: - id-token: write - contents: read target-determination: name: before-test uses: ./.github/workflows/target_determination.yml needs: llm-td - permissions: - id-token: write - contents: read get-label-type: name: get-label-type @@ -107,7 +102,6 @@ jobs: { config: "distributed", shard: 1, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.12xlarge.nvidia.gpu" }, { config: "distributed", shard: 2, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.12xlarge.nvidia.gpu" }, { config: "distributed", shard: 3, num_shards: 3, runner: "${{ needs.get-label-type.outputs.label-type }}linux.g4dn.12xlarge.nvidia.gpu" }, - { config: "pr_time_benchmarks", shard: 1, num_shards: 1, runner: "linux.g4dn.metal.nvidia.gpu" }, { config: "libtorch_agnostic_targetting", shard: 1, num_shards: 1, runner: "linux.g4dn.metal.nvidia.gpu" }, ]} secrets: inherit @@ -268,14 +262,14 @@ jobs: secrets: inherit linux-jammy-rocm-py3_10-build: - name: linux-jammy-rocm-py3.10 + name: linux-jammy-rocm-py3.10-mi355 uses: ./.github/workflows/_linux-build.yml needs: - get-label-type - job-filter with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-rocm-py3.10 + build-environment: linux-jammy-rocm-py3.10-mi355 docker-image-name: ci-image:pytorch-linux-jammy-rocm-n-py3 sync-tag: rocm-build test-matrix: | @@ -293,10 +287,7 @@ jobs: secrets: inherit linux-jammy-rocm-py3_10-test: - permissions: - id-token: write - contents: read - name: linux-jammy-rocm-py3.10 + name: linux-jammy-rocm-py3.10-mi355 uses: ./.github/workflows/_rocm-test.yml needs: - linux-jammy-rocm-py3_10-build @@ -374,9 +365,9 @@ jobs: tests-to-include: ${{ github.event.inputs.tests-to-include || '' }} secrets: inherit - linux-jammy-py3-clang15-executorch-build: -# if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3-clang15-executorch ') }} - name: linux-jammy-py3-clang15-executorch + linux-jammy-py3-clang18-executorch-build: +# if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3-clang18-executorch ') }} + name: linux-jammy-py3-clang18-executorch uses: ./.github/workflows/_linux-build.yml needs: - get-label-type @@ -384,25 +375,25 @@ jobs: if: false # Has been broken for a while with: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" - build-environment: linux-jammy-py3-clang15-executorch - docker-image-name: ci-image:pytorch-linux-jammy-py3-clang15-executorch + build-environment: linux-jammy-py3-clang18-executorch + docker-image-name: ci-image:pytorch-linux-jammy-py3-clang18-executorch test-matrix: | { include: [ { config: "executorch", shard: 1, num_shards: 1, runner: "${{ needs.get-label-type.outputs.label-type }}linux.2xlarge" }, ]} secrets: inherit - linux-jammy-py3-clang15-executorch-test: - if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3-clang15-executorch ') }} - name: linux-jammy-py3-clang15-executorch + linux-jammy-py3-clang18-executorch-test: + if: ${{ needs.job-filter.outputs.jobs == '' || contains(needs.job-filter.outputs.jobs, ' linux-jammy-py3-clang18-executorch ') }} + name: linux-jammy-py3-clang18-executorch uses: ./.github/workflows/_linux-test.yml needs: - - linux-jammy-py3-clang15-executorch-build + - linux-jammy-py3-clang18-executorch-build - job-filter with: - build-environment: ${{ needs.linux-jammy-py3-clang15-executorch-build.outputs.build-environment }} - docker-image: ${{ needs.linux-jammy-py3-clang15-executorch-build.outputs.docker-image }} - test-matrix: ${{ needs.linux-jammy-py3-clang15-executorch-build.outputs.test-matrix }} + build-environment: ${{ needs.linux-jammy-py3-clang18-executorch-build.outputs.build-environment }} + docker-image: ${{ needs.linux-jammy-py3-clang18-executorch-build.outputs.docker-image }} + test-matrix: ${{ needs.linux-jammy-py3-clang18-executorch-build.outputs.test-matrix }} tests-to-include: ${{ github.event.inputs.tests-to-include || '' }} secrets: inherit @@ -417,5 +408,35 @@ jobs: runner_prefix: "${{ needs.get-label-type.outputs.label-type }}" runner: linux.r7i.2xlarge build-environment: linux-jammy-py3.10-gcc11-full-debug-build-only - docker-image-name: ci-image:pytorch-linux-jammy-py3.10-gcc11 + docker-image-name: ci-image:pytorch-linux-jammy-py3.10-clang18 + secrets: inherit + + linux-jammy-aarch64-py3_10-gcc13-build: + name: linux-jammy-aarch64-py3.10 + uses: ./.github/workflows/_linux-build.yml + needs: get-label-type + with: + runner_prefix: ${{ needs.get-label-type.outputs.label-type }} + build-environment: linux-jammy-aarch64-py3.10 + docker-image-name: ci-image:pytorch-linux-jammy-aarch64-py3.10-gcc13 + runner: linux.arm64.m7g.4xlarge + # Periodic AArch64 tests for SVE256 coverage + test-matrix: | + { include: [ + { config: "default", shard: 1, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" }, + { config: "default", shard: 2, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" }, + { config: "default", shard: 3, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" }, + { config: "default", shard: 4, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" }, + { config: "default", shard: 5, num_shards: 5, runner: "${{ needs.get-label-type.outputs.label-type }}linux.arm64.m7g.4xlarge" }, + ]} + secrets: inherit + + linux-jammy-aarch64-py3_10-gcc13-test: + name: linux-jammy-aarch64-py3.10 + uses: ./.github/workflows/_linux-test.yml + needs: linux-jammy-aarch64-py3_10-gcc13-build + with: + build-environment: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.build-environment }} + docker-image: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.docker-image }} + test-matrix: ${{ needs.linux-jammy-aarch64-py3_10-gcc13-build.outputs.test-matrix }} secrets: inherit diff --git a/.github/workflows/update-viablestrict.yml b/.github/workflows/update-viablestrict.yml index 1b4af0f274913..439a2d8c6457e 100644 --- a/.github/workflows/update-viablestrict.yml +++ b/.github/workflows/update-viablestrict.yml @@ -23,7 +23,7 @@ jobs: with: repository: pytorch/pytorch stable-branch: viable/strict - requires: '[\"pull\", \"trunk\", \"lint\", \"linux-aarch64\"]' + requires: '[\"pull\", \"trunk\", \"lint\"]' secret-bot-token: ${{ secrets.MERGEBOT_TOKEN }} clickhouse-url: ${{ secrets.CLICKHOUSE_URL }} clickhouse-username: ${{ secrets.CLICKHOUSE_VIABLESTRICT_USERNAME }} diff --git a/.github/workflows/upload-test-stats-while-running.yml b/.github/workflows/upload-test-stats-while-running.yml index 9aecaad0e068f..7771278448a49 100644 --- a/.github/workflows/upload-test-stats-while-running.yml +++ b/.github/workflows/upload-test-stats-while-running.yml @@ -15,15 +15,11 @@ jobs: name: Upload test stats while running runs-on: linux.2xlarge steps: - - name: Checkout PyTorch - uses: pytorch/pytorch/.github/actions/checkout-pytorch@main + - name: Setup Linux + uses: pytorch/pytorch/.github/actions/setup-linux@main with: - fetch-depth: 1 submodules: false - - name: Setup Linux - uses: ./.github/actions/setup-linux - - name: Install requirements run: | python3 -m pip install requests==2.32.2 boto3==1.35.42 diff --git a/.github/workflows/upload-test-stats.yml b/.github/workflows/upload-test-stats.yml index 5dcbfe7fd65fa..8d274071662e1 100644 --- a/.github/workflows/upload-test-stats.yml +++ b/.github/workflows/upload-test-stats.yml @@ -20,6 +20,7 @@ on: - rocm-mi300 - rocm-mi355 - rocm-navi31 + - rocm-nightly - inductor-micro-benchmark - inductor-micro-benchmark-x86 - inductor-cu124 @@ -29,7 +30,6 @@ on: - inductor-perf-nightly-rocm-mi300 - inductor-perf-nightly-rocm-mi355 - mac-mps - - linux-aarch64 types: - completed diff --git a/.github/workflows/vllm.yml b/.github/workflows/vllm.yml index eb1c78019b10e..f65e73042c245 100644 --- a/.github/workflows/vllm.yml +++ b/.github/workflows/vllm.yml @@ -16,6 +16,7 @@ concurrency: permissions: id-token: write contents: read + actions: read jobs: build: diff --git a/.github/workflows/win-arm64-build-test.yml b/.github/workflows/win-arm64-build-test.yml index 95b4e2f027f60..5c31628b96871 100644 --- a/.github/workflows/win-arm64-build-test.yml +++ b/.github/workflows/win-arm64-build-test.yml @@ -31,7 +31,7 @@ jobs: steps: - name: configure aws credentials id: aws_creds - uses: aws-actions/configure-aws-credentials@v4 + uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4 with: role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_sscache aws-region: us-east-1 diff --git a/.github/workflows/xpu.yml b/.github/workflows/xpu.yml index fbfa9d62cc571..8e089db1f9cee 100644 --- a/.github/workflows/xpu.yml +++ b/.github/workflows/xpu.yml @@ -14,6 +14,11 @@ concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true +permissions: + id-token: write + contents: read + actions: read + jobs: get-label-type: @@ -78,9 +83,6 @@ jobs: name: linux-noble-xpu-n-py3.10 uses: ./.github/workflows/_xpu-test.yml needs: linux-noble-xpu-n-py3_10-build - permissions: - id-token: write - contents: read with: build-environment: ${{ needs.linux-noble-xpu-n-py3_10-build.outputs.build-environment }} docker-image: ${{ needs.linux-noble-xpu-n-py3_10-build.outputs.docker-image }} @@ -107,9 +109,6 @@ jobs: name: linux-noble-xpu-n-py3.10-client uses: ./.github/workflows/_xpu-test.yml needs: linux-noble-xpu-n-py3_10-client-build - permissions: - id-token: write - contents: read with: build-environment: ${{ needs.linux-noble-xpu-n-py3_10-client-build.outputs.build-environment }} docker-image: ${{ needs.linux-noble-xpu-n-py3_10-client-build.outputs.docker-image }} diff --git a/.gitignore b/.gitignore index 54e3fa32b2fdb..5105977ae52c0 100644 --- a/.gitignore +++ b/.gitignore @@ -8,8 +8,7 @@ ## PyTorch -.coverage -coverage.xml +agent_space/ .dmypy.json .gradle .hypothesis @@ -40,7 +39,6 @@ docs/build/ docs/cpp/src docs/src/**/* docs/cpp/build -docs/cpp/source/api docs/cpp/source/html/ docs/cpp/source/latex/ docs/source/compile/generated/ @@ -52,7 +50,6 @@ usage_log* test-reports/ test/*.bak test/**/*.bak -test/.coverage test/.hypothesis/ test/cpp/api/mnist test/custom_operator/model.pt @@ -367,7 +364,13 @@ xla/ pr.diff # coverage files -*/**/.coverage.* +.coverage.* +.coverage +coverage.xml +*.profdata +*.profraw +coverage-html +coverage.lcov # buck generated files .buckd/ diff --git a/.lintrunner.toml b/.lintrunner.toml index 5c42dc81c44d8..9849094fdda9e 100644 --- a/.lintrunner.toml +++ b/.lintrunner.toml @@ -191,6 +191,11 @@ exclude_patterns = [ 'aten/src/ATen/cpu/FlushDenormal.cpp', 'aten/src/ATen/cpu/vml.h', 'aten/src/ATen/CPUFixedAllocator.h', + # interned_strings.h includes the generated aten_interned_strings.h which + # #errors under TORCH_ASSERT_ONLY_METHOD_OPERATORS. The .cpp file #undefs + # the macro before including the header, but clang-tidy lints the .h + # standalone using the .cpp's compile flags, so the #error always fires. + 'aten/src/ATen/core/interned_strings.h', 'aten/src/ATen/Parallel*.h', 'c10/xpu/**/*.h', 'c10/xpu/**/*.cpp', @@ -532,6 +537,8 @@ exclude_patterns = [ 'torch/csrc/inductor/aoti_runtime/**', # Test files use EXPECT_THROW which is a gtest macro 'test/cpp/**/*.cpp', + # JIT frontend uses throw(ErrorReport(...)) pervasively + 'torch/csrc/jit/frontend/schema_type_parser.cpp', ] command = [ 'python3', diff --git a/.spin/cmds.py b/.spin/cmds.py index f31acfa977424..71401b14e9729 100644 --- a/.spin/cmds.py +++ b/.spin/cmds.py @@ -8,6 +8,11 @@ import spin +CWD = Path(__file__).absolute().parent.parent +sys.path.insert(0, str(CWD)) # this only affects the current process +from tools.clean import clean as _clean + + def file_digest(file, algorithm: str): try: return hashlib.file_digest(file, algorithm) @@ -441,6 +446,12 @@ def quickfix(ctx, *, lintrunner_args, **kwargs): ctx.invoke(quicklint, apply_patches=True) +@click.command() +def clean(): + """Clean, that is remove all files in .gitignore except in the NOT-CLEAN-FILES section.""" + _clean() + + @click.command() def regenerate_github_workflows(): """Regenerate GitHub workflows from templates.""" diff --git a/AGENTS.md b/AGENTS.md deleted file mode 100644 index 718217d3e663d..0000000000000 --- a/AGENTS.md +++ /dev/null @@ -1,18 +0,0 @@ -- This is the only AGENTS.md, there are no recursive AGENTS.md -- When you are working on a bug, first create a standalone file that - reproduces the bug and verify it fails in the expected way. Use this to - test if your changes work. Once the change is passing, find an appropriate - test file to add the test to and make sure to follow local conventions on - the test file. -- If you are running the real test suite, DO NOT run the entire test suite. - Instead run only a single test case, e.g., 'python test/test_torch.py TestTorch.test_dir' -- Do NOT run setup.py, you do not have a working build environment -- Do NOT run pre-commit, it is not setup -- To run lint, run 'lintrunner -a' (which will autoapply changes) -- Do NOT attempt to install dependencies, you do not have Internet access -- Do NOT create summary files unless explicitly asked -- When you are ready to make a PR, do exactly these steps: - - git stash -u - - git reset --hard $(cat /tmp/orig_work.txt) # NB: reset to the LOCAL branch, do NOT fetch - - git stash pop - - Resolve conflicts if necessary diff --git a/AGENTS.md b/AGENTS.md new file mode 120000 index 0000000000000..681311eb9cf45 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1 @@ +CLAUDE.md \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md index e83bcdf5e0e12..640a24bbcb410 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,3 +1,7 @@ +# Scratch Space + +Use `agent_space/` (git-ignored, at repo root) for temporary scripts, scratch files, and throwaway experiments. Do not commit files from this directory. + # PR Review When asked to review a PR, always use the /pr-review skill. diff --git a/CMakeLists.txt b/CMakeLists.txt index fdb5062824815..d5941243d66db 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -27,6 +27,17 @@ endif() # ---[ Project and semantic versioning. project(Torch CXX C) +# When built via scikit-build-core, Python_EXECUTABLE is already set. +# For standalone cmake invocations, find it via find_package. +find_package(Python COMPONENTS Interpreter REQUIRED) + +# Forward environment variables to CMake variables, replicating the behavior +# previously handled by setup.py / tools/setup_helpers/cmake.py. +# When built via setup.py, these are already set as -D flags; the module +# skips variables that are already defined, so it is safe to include +# unconditionally. +include(cmake/EnvVarForwarding.cmake) + if(${CMAKE_SYSTEM_NAME} STREQUAL "Linux") set(LINUX TRUE) else() @@ -249,6 +260,7 @@ endif() option(USE_CUDA "Use CUDA" ON) option(USE_XPU "Use XPU" ON) +option(USE_MTIA "Use MTIA" OFF) cmake_dependent_option( BUILD_LAZY_CUDA_LINALG "Build cuda linalg ops as separate library" ON "USE_CUDA AND LINUX AND BUILD_PYTHON" OFF) @@ -291,7 +303,7 @@ option(USE_NNPACK "Use NNPACK" ON) cmake_dependent_option(USE_NUMA "Use NUMA. Only available on Linux." ON "LINUX" OFF) cmake_dependent_option(USE_NVRTC "Use NVRTC. Only available if USE_CUDA is on." - OFF "USE_CUDA" OFF) + ON "USE_CUDA" OFF) option(USE_NUMPY "Use NumPy" ON) option(USE_OBSERVERS "Use observers module." OFF) option(USE_OPENCL "Use OpenCL" OFF) @@ -590,7 +602,7 @@ if(MSVC) # https://learn.microsoft.com/en-us/cpp/build/reference/zc-cplusplus?view=msvc-170 # * https://en.cppreference.com/w/cpp/preprocessor/replace#Predefined_macros set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /Zc:__cplusplus") - set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler /Zc:__cplusplus") + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler /Zc:__cplusplus -Xcompiler /Zc:preprocessor") set(CMAKE_NINJA_CMCLDEPS_RC OFF) if(MSVC_Z7_OVERRIDE) @@ -1181,6 +1193,12 @@ if(NOT MSVC) string(APPEND CMAKE_CXX_FLAGS_DEBUG " -fno-omit-frame-pointer -O0") string(APPEND CMAKE_LINKER_FLAGS_DEBUG " -fno-omit-frame-pointer -O0") endif() + # aarch64 C++ stack unwinding uses frame-pointer chain walking, so frame + # pointers must be present in all build types. The cost is negligible on + # aarch64 (31 GPRs vs x86-64's 16, so dedicating x29 rarely spills). + if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") + append_cxx_flag_if_supported("-fno-omit-frame-pointer" CMAKE_CXX_FLAGS) + endif() append_cxx_flag_if_supported("-fno-math-errno" CMAKE_CXX_FLAGS) append_cxx_flag_if_supported("-fno-trapping-math" CMAKE_CXX_FLAGS) append_cxx_flag_if_supported("-Werror=format" CMAKE_CXX_FLAGS) @@ -1229,10 +1247,11 @@ if(USE_CPP_CODE_COVERAGE) if("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU") string(APPEND CMAKE_C_FLAGS " --coverage -fprofile-abs-path") string(APPEND CMAKE_CXX_FLAGS " --coverage -fprofile-abs-path") + string(APPEND CMAKE_OBJCXX_FLAGS " --coverage -fprofile-abs-path") elseif("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang") string(APPEND CMAKE_C_FLAGS " -fprofile-instr-generate -fcoverage-mapping") - string(APPEND CMAKE_CXX_FLAGS - " -fprofile-instr-generate -fcoverage-mapping") + string(APPEND CMAKE_CXX_FLAGS " -fprofile-instr-generate -fcoverage-mapping") + string(APPEND CMAKE_OBJCXX_FLAGS " -fprofile-instr-generate -fcoverage-mapping") else() message( ERROR diff --git a/CODEOWNERS b/CODEOWNERS index 5fc084f4799f1..604ad3c694d9a 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -48,6 +48,17 @@ nn/qat/ @jerryzh168 # c10d backend APIs /torch/csrc/distributed/c10d/Backend.* @kwen2501 /torch/csrc/distributed/c10d/Ops.* @kwen2501 +# FSDP, DDP +/torch/distributed/_composable/fsdp/ @weifengpy +/torch/distributed/_composable/replicate.py @weifengpy +/torch/distributed/_composable/replicate_with_fsdp.py @weifengpy +/torch/distributed/fsdp/ @weifengpy +/torch/distributed/optim/ @weifengpy +/torch/nn/parallel/ @weifengpy +# DTensor view, matmul, _StridedShard +/torch/distributed/tensor/placement_types.py @weifengpy +/torch/distributed/tensor/_ops/_view_ops.py @weifengpy +/torch/distributed/tensor/_ops/_matrix_ops.py @weifengpy # ONNX Export /torch/_dynamo/backends/onnxrt.py @titaiwangms @xadupre @justinchuby @@ -173,9 +184,9 @@ caffe2/utils/hip @jeffdaily @jithunnair-amd /torch/_export/serde/schema.py @SherlockNoMad @zhxchen17 # Dynamic Shapes -/torch/fx/experimental/symbolic_shapes.py @bobrenjc93 @laithsakka -/torch/fx/experimental/sym_node.py @bobrenjc93 @laithsakka -/torch/fx/experimental/recording.py @bobrenjc93 @laithsakka +/torch/fx/experimental/symbolic_shapes.py @laithsakka +/torch/fx/experimental/sym_node.py @laithsakka +/torch/fx/experimental/recording.py @laithsakka # ProxyTorchDispatchMode torch/fx/experimental/proxy_tensor.py @aorenste diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 4e5850e354814..e92aec464946a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -16,6 +16,7 @@ aspects of contributing to PyTorch. - [Codebase structure](#codebase-structure) - [AI-Assisted Development](#ai-assisted-development) - [Spin](#spin) + - [Building](#building) - [Linting](#linting) - [default lint](#default-lint) - [Regenerating](#regenerating) @@ -300,6 +301,14 @@ helps running common tasks. To list the available tasks, run `spin --help`. Currently, we support the following tasks with Spin: +### Building + +To support building and general development, the following commands exist. + +|command|| +|-|-| +|`clean`|clean, that is remove files and directories listed in .gitignore before the NOT-CLEAN-FILES marker| + ### Linting Spin helps with linting by making sure that lintrunner is installed correctly diff --git a/RELEASE.md b/RELEASE.md index 6d10e359c1324..3120931b38557 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -130,7 +130,7 @@ Following requirements need to be met prior to cutting a release branch: * Triton release branch must be created (e.g., [release/3.6.x](https://github.com/triton-lang/triton/tree/release/3.6.x)) and the Triton pin update PR must be landed (e.g., [#168096](https://github.com/pytorch/pytorch/pull/168096)) at least 1 week before the branch cut * Resolve all outstanding issues in the milestones that are feature work and release blocking (for example [release 2.10 milestone](https://github.com/pytorch/pytorch/milestone/57)). A report of outstanding cherry-picks can be produced by running the [github-analytics-daily workflow](https://github.com/pytorch/test-infra/blob/main/.github/workflows/github-analytics-daily.yml) * Validate that all new workflows have been created in the PyTorch and domain libraries included in the release. Validate it against all dimensions of release matrix, including operating systems (Linux, macOS, Windows), Python versions as well as CPU architectures (x86 and arm) and accelerator versions (CUDA, ROCm, XPU). -* All [viable/strict](.github/workflows/update-viablestrict.yml) jobs are green, which requires the following jobs to pass: `pull`, `trunk`, `lint`, `linux-aarch64` +* All [viable/strict](.github/workflows/update-viablestrict.yml) jobs are green, which requires the following jobs to pass: `pull`, `trunk`, `lint` * All the nightly jobs for pytorch and domain libraries should be green. Validate this using the following HUD links: * [PyTorch](https://hud.pytorch.org/hud/pytorch/pytorch/nightly) * [TorchVision](https://hud.pytorch.org/hud/pytorch/vision/nightly) diff --git a/WORKSPACE b/WORKSPACE index 6792a747ecbc5..be911d65c5580 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -210,11 +210,8 @@ http_archive( http_archive( name = "rules_python", - # TODO Fix bazel linter to support hashes for release tarballs. - # - # sha256 = "94750828b18044533e98a129003b6a68001204038dc4749f40b195b24c38f49f", - strip_prefix = "rules_python-0.21.0", - url = "https://github.com/bazelbuild/rules_python/releases/download/0.21.0/rules_python-0.21.0.tar.gz", + strip_prefix = "rules_python-0.31.0", + url = "https://github.com/bazelbuild/rules_python/releases/download/0.31.0/rules_python-0.31.0.tar.gz", ) load("@rules_python//python:repositories.bzl", "py_repositories") @@ -236,6 +233,13 @@ pip_parse( name = "pip_deps", python_interpreter_target = interpreter, requirements_lock = "//:tools/build/bazel/requirements.txt", + extra_pip_args = [ + "--index-url", + "${PIP_INDEX_URL}", + "--trusted-host", + "${PIP_TRUSTED_HOST}", + ], + envsubst = ["PIP_INDEX_URL", "PIP_TRUSTED_HOST"], ) load("@pip_deps//:requirements.bzl", "install_deps") diff --git a/android/pytorch_android/CMakeLists.txt b/android/pytorch_android/CMakeLists.txt index c640cbdeac9f1..4ee3011683011 100644 --- a/android/pytorch_android/CMakeLists.txt +++ b/android/pytorch_android/CMakeLists.txt @@ -14,7 +14,7 @@ endif() include(GNUInstallDirs) -set(CMAKE_CXX_STANDARD 17 CACHE STRING "The C++ standard whose features are requested to build this target.") +set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.") set(CMAKE_VERBOSE_MAKEFILE ON) message(STATUS "ANDROID_STL:${ANDROID_STL}") diff --git a/android/pytorch_android_torchvision/CMakeLists.txt b/android/pytorch_android_torchvision/CMakeLists.txt index 2c8931f3fb911..b49b779a18b1d 100644 --- a/android/pytorch_android_torchvision/CMakeLists.txt +++ b/android/pytorch_android_torchvision/CMakeLists.txt @@ -1,6 +1,6 @@ cmake_minimum_required(VERSION 3.10) project(pytorch_vision_jni CXX) -set(CMAKE_CXX_STANDARD 17 CACHE STRING "The C++ standard whose features are requested to build this target.") +set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.") set(CMAKE_VERBOSE_MAKEFILE ON) set(pytorch_vision_cpp_DIR ${CMAKE_CURRENT_LIST_DIR}/src/main/cpp) diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt index 1e06d30774d90..06522a9298893 100644 --- a/aten/src/ATen/CMakeLists.txt +++ b/aten/src/ATen/CMakeLists.txt @@ -892,12 +892,23 @@ if(USE_MPS) metal_to_air(${SHADER} ${TGT_BASIC} "-std=metal3.1") endforeach() air_to_metallib(kernels_basic.metallib ${AIR_BASIC}) + set(METALLIB_DEPS kernels_basic.metallib) + if(CAN_COMPILE_METAL_40) + foreach(SHADER ${native_mps_metal} ${native_mps_sparse_metal}) + cmake_path(GET SHADER STEM TGT_STEM) + string(CONCAT TGT_40 ${TGT_STEM} "_40.air") + list(APPEND AIR_40 ${TGT_40}) + metal_to_air(${SHADER} ${TGT_40} "-std=metal4.0") + endforeach() + air_to_metallib(kernels_40.metallib ${AIR_40}) + list(APPEND METALLIB_DEPS kernels_40.metallib) + endif() add_custom_command( COMMAND echo "// $$(date)" > metallib_dummy.cpp - DEPENDS kernels_basic.metallib + DEPENDS ${METALLIB_DEPS} OUTPUT metallib_dummy.cpp COMMENT "Updating metallibs timestamp") - add_custom_target(metallibs DEPENDS kernels_basic.metallib metallib_dummy.cpp) + add_custom_target(metallibs DEPENDS ${METALLIB_DEPS} metallib_dummy.cpp) else() file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/native/mps") foreach(SHADER ${native_mps_metal} ${native_mps_sparse_metal}) @@ -925,7 +936,7 @@ install(FILES "${CMAKE_CURRENT_BINARY_DIR}/cmake-exports/ATenConfig.cmake" set(INSTALL_HEADERS ${base_h} ${ATen_CORE_HEADERS} ${native_nested_h} ${ATen_TRANSFORMER_HEADERS}) if(NOT INTERN_BUILD_MOBILE) - list(APPEND INSTALL_HEADERS ${native_h} ${native_cpu_h} ${native_ao_sparse_h} ${native_quantized_h} ${cuda_h} ${native_cuda_h} ${native_hip_h} ${native_mtia_h} ${cudnn_h} ${hip_h} ${mtia_h} ${xpu_h} ${mps_h} ${native_kleidiai_h} ${native_mps_h} ${native_utils_h} ${miopen_h} ${mkldnn_xpu_h}) + list(APPEND INSTALL_HEADERS ${native_h} ${native_cpu_h} ${native_ao_sparse_h} ${native_quantized_h} ${cuda_h} ${native_cuda_h} ${native_hip_h} ${native_mtia_h} ${cudnn_h} ${hip_h} ${mtia_h} ${xpu_h} ${mps_h} ${native_kleidiai_h} ${native_mps_h} ${native_utils_h} ${miopen_h}) # Metal if(USE_PYTORCH_METAL_EXPORT) # Add files needed from exporting metal models(optimized_for_mobile) diff --git a/aten/src/ATen/DeviceAccelerator.cpp b/aten/src/ATen/DeviceAccelerator.cpp index efab9ec9c5927..fed2e6c4febe8 100644 --- a/aten/src/ATen/DeviceAccelerator.cpp +++ b/aten/src/ATen/DeviceAccelerator.cpp @@ -1,4 +1,5 @@ #include +#include #include #include @@ -136,6 +137,11 @@ c10::DeviceCapability getDeviceCapability(c10::DeviceIndex device_index) { c10::impl::VirtualGuardImpl impl(device_type); return impl.getDeviceCapability({device_type, device_index}); } + +void emptyHostCache() { + const auto device_type = getAccelerator(true).value(); + at::getHostAllocator(device_type)->empty_cache(); +} // NOLINTEND(bugprone-unchecked-optional-access) } // namespace at::accelerator diff --git a/aten/src/ATen/DeviceAccelerator.h b/aten/src/ATen/DeviceAccelerator.h index 81678ba6efc29..90fad38a5faca 100644 --- a/aten/src/ATen/DeviceAccelerator.h +++ b/aten/src/ATen/DeviceAccelerator.h @@ -78,11 +78,19 @@ TORCH_API c10::DeviceIndex maybeExchangeDevice(c10::DeviceIndex device_index); TORCH_API c10::DeviceCapability getDeviceCapability( c10::DeviceIndex device_index); +// Releases all unused device memory currently held by the accelerator's +// device-side caching allocator. The freed memory becomes available for reuse +// by other applications or processes. TORCH_API inline void emptyCache() { const auto device_type = getAccelerator(true).value(); at::getDeviceAllocator(device_type)->emptyCache(); } +// Releases all unused host (pinned) memory currently held by the accelerator's +// host-side caching allocator. The freed memory becomes available for reuse by +// other applications or processes. +TORCH_API void emptyHostCache(); + TORCH_API inline at::CachingDeviceAllocator::DeviceStats getDeviceStats( c10::DeviceIndex device_index) { const auto device_type = getAccelerator(true).value(); diff --git a/aten/src/ATen/core/Vitals.cpp b/aten/src/ATen/core/Vitals.cpp deleted file mode 100644 index db58c03830539..0000000000000 --- a/aten/src/ATen/core/Vitals.cpp +++ /dev/null @@ -1,96 +0,0 @@ -#include -#include -#include - -namespace at::vitals { - -APIVitals VitalsAPI; - -std::ostream& operator<<(std::ostream& os, TorchVital const& tv) { - for (const auto& m : tv.attrs) { - os << "[TORCH_VITAL] " << tv.name << '.' << m.first << "\t\t " - << m.second.value << '\n'; - } - return os; -} - -TorchVital::~TorchVital() { - if (torchVitalEnabled()) { - std::cout << *this; - } -} - -TorchVitalAttr& TorchVital::create(const std::string& attr) { - return create(attr, /* force = */ false); -} - -TorchVitalAttr& TorchVital::create(const std::string& attr, bool force) { - if (!(torchVitalEnabled() || force)) { - static TorchVitalAttr disabled; - return disabled; - } - auto iter = attrs.find(attr); - if (iter == attrs.end()) { - auto r = attrs.emplace(attr, TorchVitalAttr()); - return r.first->second; - } - return iter->second; -} - -bool torchVitalEnabled() { - // If this is a performance hit, make `enabled` variable static - // and return `const bool&` instead - bool enabled = []() { - auto const e = c10::utils::get_env("TORCH_VITAL"); - if (e.has_value()) { - return !e.value().empty(); - } - return false; - }(); - if (enabled) { - VitalsAPI.vitals_enabled = true; - } - return VitalsAPI.vitals_enabled; -} - -std::string APIVitals::readVitals() { - if (!torchVitalEnabled()) { - return ""; - } - - std::stringstream buf; - for (const auto& x : name_map_) { - buf << x.second; - } - return buf.str(); -} - -bool APIVitals::setVital( - const std::string& vital_name, - const std::string& attr_name, - const std::string& value, - bool force) { - if (!(torchVitalEnabled() || force)) { - return false; - } - - auto iter = name_map_.find(vital_name); - TorchVital* vital = nullptr; - if (iter == name_map_.end()) { - auto r = name_map_.emplace(vital_name, TorchVital(vital_name)); - vital = &r.first->second; - } else { - vital = &iter->second; - } - - vital->create(attr_name, force).write(value, force); - return true; -} - -APIVitals::APIVitals() : vitals_enabled(false) { - // Set default values, force is necessary because in unit tests the env - // variable may not be set when global APIVitals are constructed. - setVital("CUDA", "used", "False", /* force = */ true); -} - -} // namespace at::vitals diff --git a/aten/src/ATen/core/Vitals.h b/aten/src/ATen/core/Vitals.h deleted file mode 100644 index 2fd7729744a10..0000000000000 --- a/aten/src/ATen/core/Vitals.h +++ /dev/null @@ -1,94 +0,0 @@ -#pragma once -#include -#include -#include - -#include - -namespace at::vitals { - -TORCH_API bool torchVitalEnabled(); - -struct TORCH_API TorchVitalAttr { - // always initialized to empty - std::string value; - template - TorchVitalAttr& operator<<(const T& t) { - if (torchVitalEnabled()) { - std::stringstream ss; - ss << t; - value += ss.str(); - } - return *this; - } - - template - void write(const T& t, bool force) { - if (force || torchVitalEnabled()) { - std::stringstream ss; - ss << t; - value = ss.str(); - } - } -}; - -struct TORCH_API TorchVital { - std::string name; - std::unordered_map attrs; - - explicit TorchVital(std::string n) : name(std::move(n)) {} - TorchVital(const TorchVital&) = default; - TorchVital(TorchVital&&) = default; - TorchVital& operator=(const TorchVital&) = default; - TorchVital& operator=(TorchVital&&) = default; - TorchVital() = delete; - - TorchVitalAttr& create(const std::string& attr); - TorchVitalAttr& create(const std::string& attr, bool force); - friend std::ostream& operator<<(std::ostream& os, const TorchVital& dt); - - ~TorchVital(); -}; - -std::ostream& operator<<(std::ostream& os, TorchVital const& tv); - -// A way to access vitals by string names instead of by global reference. -// This enables access to vitals from the PythonAPI. -class TORCH_API APIVitals { - public: - bool vitals_enabled; - - // Set any vital sign that was added to the map. - bool setVital( - const std::string& vital_name, - const std::string& attr_name, - const std::string& value, - bool force = false); - std::string readVitals(); - - APIVitals(); - - // Ensure this stays a singleton - APIVitals(APIVitals const& other) = delete; - APIVitals(APIVitals&& other) = delete; - APIVitals& operator=(const APIVitals&) = delete; - APIVitals& operator=(APIVitals&&) = delete; - ~APIVitals() = default; - - private: - std::unordered_map name_map_; -}; - -extern TORCH_API APIVitals VitalsAPI; - -} // namespace at::vitals - -#define TORCH_VITAL_DECLARE(name) \ - TORCH_API at::vitals::TorchVital TorchVital_##name; - -#define TORCH_VITAL_DEFINE(name) \ - TORCH_API at::vitals::TorchVital TorchVital_##name(#name); - -#define TORCH_VITAL_BASE(name) TorchVital_##name - -#define TORCH_VITAL(name, attr) TORCH_VITAL_BASE(name).create(#attr) diff --git a/aten/src/ATen/core/interned_strings.h b/aten/src/ATen/core/interned_strings.h index 38942031befcd..532c921ca4258 100644 --- a/aten/src/ATen/core/interned_strings.h +++ b/aten/src/ATen/core/interned_strings.h @@ -187,7 +187,6 @@ namespace c10 { _(aten, append) \ _(aten, as_tensor) \ _(aten, adaptive_avg_pool2d_backward) \ - _(aten, dim) \ _(aten, format) \ _(aten, percentFormat) \ _(aten, __not__) \ diff --git a/aten/src/ATen/cpu/vec/vec512/vec512_float8.h b/aten/src/ATen/cpu/vec/vec512/vec512_float8.h index 2a4ca9b958c55..60035b7bee75c 100644 --- a/aten/src/ATen/cpu/vec/vec512/vec512_float8.h +++ b/aten/src/ATen/cpu/vec/vec512/vec512_float8.h @@ -377,6 +377,15 @@ class Vectorizedf8 { } } + static Vectorized blendv( + const Vectorized& a, + const Vectorized& b, + const Vectorized& mask) { + auto msb_one = _mm512_set1_epi8(0xFF); + auto mask_ = _mm512_cmp_epu8_mask((__m512i)mask, msb_one, _MM_CMPINT_EQ); + return _mm512_mask_blend_epi8(mask_, (__m512i)a, (__m512i)b); + } + Vectorized abs() const { return _mm512_andnot_si512(_mm512_set1_epi8(0x80), values); } diff --git a/aten/src/ATen/cpu/vec/vec_convert.h b/aten/src/ATen/cpu/vec/vec_convert.h index 5f43b475f0f89..cce9d46910e78 100644 --- a/aten/src/ATen/cpu/vec/vec_convert.h +++ b/aten/src/ATen/cpu/vec/vec_convert.h @@ -63,12 +63,24 @@ inline std::enable_if_t, Vectorized> convert return src; } +template +inline std::enable_if_t, Vectorized> +round_convert(const Vectorized& src) { + return src; +} + template inline std::enable_if_t, Vectorized> convert(const Vectorized& src) { return VecConvert::apply(src); } +template +inline std::enable_if_t, Vectorized> +round_convert(const Vectorized& src) { + return VecRoundConvert::apply(src); +} + template < typename dst_t, int dst_n, @@ -90,12 +102,6 @@ inline VectorizedN round_convert( return VecRoundConvert::apply(src); } -template -inline std::enable_if_t, Vectorized> -round_convert(const Vectorized& src) { - return VecRoundConvert::apply(src); -} - template < typename dst_t, int dst_n, @@ -108,6 +114,18 @@ convert(const VectorizedN& src) { return VecConvert::apply(src); } +template < + typename dst_t, + int dst_n, + typename src_t, + int src_n, + bool keep = false, + std::enable_if_t = 0> +inline std::conditional_t, Vectorized> +round_convert(const VectorizedN& src) { + return VecRoundConvert::apply(src); +} + } // namespace CPU_CAPABILITY template < diff --git a/aten/src/ATen/cuda/CUDAContextLight.h b/aten/src/ATen/cuda/CUDAContextLight.h index 6251d40e6daad..e93d5b19d56bf 100644 --- a/aten/src/ATen/cuda/CUDAContextLight.h +++ b/aten/src/ATen/cuda/CUDAContextLight.h @@ -85,7 +85,7 @@ TORCH_CUDA_CPP_API c10::Allocator* getCUDADeviceAllocator(); /* Handles */ TORCH_CUDA_CPP_API cusparseHandle_t getCurrentCUDASparseHandle(); -TORCH_CUDA_CPP_API cublasHandle_t getCurrentCUDABlasHandle(); +TORCH_CUDA_CPP_API cublasHandle_t getCurrentCUDABlasHandle(bool setup = true); TORCH_CUDA_CPP_API cublasLtHandle_t getCurrentCUDABlasLtHandle(); TORCH_CUDA_CPP_API void clearCublasWorkspaces(); diff --git a/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp b/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp index 536020ffa1c86..9b690502bf7a5 100644 --- a/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp +++ b/aten/src/ATen/cuda/CUDAGeneratorImpl.cpp @@ -82,151 +82,144 @@ Generator createCUDAGenerator(DeviceIndex device_index) { } // namespace cuda::detail /** - * Creates a clone of this CUDA Generator State. + * Allocate GPU tensors for this capture state. + * + * We allocate on the default stream so that the caching allocator routes + * these tensors to the default memory pool, not the graph's capture pool. */ +void CUDAGeneratorCaptureState::initialize(uint64_t seed) { + if (is_initialized()) { + return; + } + + auto options = at::TensorOptions().device(at::kCUDA).dtype(at::kLong); + c10::InferenceMode inference_guard(false); + + // Allocate on the default stream so that the caching allocator routes + // these tensors to the default memory pool, not the graph's capture pool. + // The relaxed capture mode guard is needed because the thread-local capture + // mode may be Global (set by cudaStreamBeginCapture), which would block + // cudaMalloc even on a non-capturing stream. + c10::cuda::CUDAStreamCaptureModeGuard capture_mode_guard( + cudaStreamCaptureModeRelaxed); + c10::cuda::CUDAStreamGuard stream_guard(c10::cuda::getDefaultCUDAStream()); + + rng_state_seed_extragraph_ = at::empty({1}, options); + rng_state_offset_extragraph_ = at::empty({1}, options); + + // Synchronize the default stream so that any prior work completes before + // a different stream writes to this memory. + c10::cuda::getDefaultCUDAStream().synchronize(); + + offset_intragraph_ = 0; +} + +void CUDAGeneratorCaptureState::increase(uint64_t increment) { + // see Note [Why enforce RNG offset % 4 == 0?] + TORCH_INTERNAL_ASSERT( + offset_intragraph_ % 4 == 0, "RNG offset must be a multiple of 4."); + TORCH_INTERNAL_ASSERT( + offset_intragraph_ <= std::numeric_limits::max() - increment, + "Increment causes overflow in the offset value."); + offset_intragraph_ += increment; +} + +uint64_t CUDAGeneratorCaptureState::finalize() { + uint64_t result = offset_intragraph_; + offset_intragraph_ = 0; + return result; +} + +void CUDAGeneratorCaptureState::setup_for_replay(uint64_t seed, uint64_t philox_offset) { + TORCH_INTERNAL_ASSERT(is_initialized(), + "Capture state not initialized"); + rng_state_seed_extragraph_.fill_(static_cast(seed)); + rng_state_offset_extragraph_.fill_(static_cast(philox_offset)); +} + c10::intrusive_ptr CUDAGeneratorState::clone() { - return make_intrusive( - seed_, philox_offset_per_thread_, offset_intragraph_); + return make_intrusive(seed_, philox_offset_per_thread_); } /** - * Function to increase the internal offset based on the specified increment. + * Lookup capture state for a capture ID. Returns nullptr if not found. */ -void CUDAGeneratorState::increase(uint64_t increment) { - // Rounds increment up to the nearest multiple of 4 to meet alignment - // requirements. - // see Note [Why enforce RNG offset % 4 == 0?] - increment = ((increment + 3) / 4) * 4; - // Handling different behaviors based on whether capturing is active. - if (at::cuda::currentStreamCaptureStatus() != at::cuda::CaptureStatus::None) { - // Ensures that the state is actually capturing. - TORCH_CHECK( - capturing_, - "Attempt to increase offset for a CUDA generator not in capture mode."); - // Ensures the offset is a multiple of 4 - // see Note [Why enforce RNG offset % 4 == 0?] - TORCH_INTERNAL_ASSERT( - offset_intragraph_ % 4 == 0, "RNG offset must be a multiple of 4."); - // Ensures the increment does not cause overflow. - TORCH_INTERNAL_ASSERT( - offset_intragraph_ <= std::numeric_limits::max() - increment, - "Increment causes overflow in the offset value."); - offset_intragraph_ += increment; - } else { - // Checks that the increment is expected outside graph capturing. - TORCH_CHECK( - !capturing_, - "Offset increment outside graph capture encountered unexpectedly."); - // Ensures the offset is a multiple of 4 - // see Note [Why enforce RNG offset % 4 == 0?] - TORCH_INTERNAL_ASSERT( - philox_offset_per_thread_ % 4 == 0, - "RNG offset must be a multiple of 4."); - philox_offset_per_thread_ += increment; +CUDAGeneratorCaptureState* CUDAGeneratorState::get_capture_state(CaptureId_t capture_id) { + std::lock_guard lock(capture_states_mutex_); + auto it = capture_states_.find(capture_id); + if (it != capture_states_.end()) { + return it->second.get(); } + return nullptr; } /** - * Registers this state to a CUDA graph to manage within the graph. + * Create and initialize capture state for a given capture ID. + * Called during capture_begin for each registered generator. */ -void CUDAGeneratorState::register_graph(cuda::CUDAGraph* graph) { - // Ensures that the RNG state is not currently being captured. - at::cuda::assertNotCapturing( - "Cannot register the state during capturing stage."); - - // If this is the first graph to be registered, allocate memory for the seed - // and offset on the GPU. - if (registered_graphs_.empty()) { - auto options = at::TensorOptions().device(at::kCUDA).dtype(at::kLong); - // Create these tensors outside of inference mode to ensure they can be - // modified in-place later. If we create them as inference tensors, - // subsequent fill_() calls outside inference mode - // will fail with "Inplace update to inference tensor outside InferenceMode". - c10::InferenceMode guard(false); - seed_extragraph_ = at::empty({1}, options); - offset_extragraph_ = at::empty({1}, options); +void CUDAGeneratorState::init_capture_state(CaptureId_t capture_id) { + { + std::lock_guard lock(capture_states_mutex_); + if (capture_states_.count(capture_id)) { + return; + } } - // Insert the graph into the set of registered graphs if it's not already - // registered. - if (registered_graphs_.find(graph) == registered_graphs_.end()) { - registered_graphs_.insert(graph); + auto capture_state = make_intrusive(); + capture_state->initialize(seed_); + + std::lock_guard lock(capture_states_mutex_); + if (!capture_states_.count(capture_id)) { + capture_states_[capture_id] = std::move(capture_state); } } /** - * Unregisters a CUDA graph from the RNG state. + * Function to increase the internal offset based on the specified increment. */ -void CUDAGeneratorState::unregister_graph(cuda::CUDAGraph* graph) { - // Verify the graph was previously registered. - TORCH_CHECK( - registered_graphs_.find(graph) != registered_graphs_.end(), - "The graph should be registered to the state"); - - // Remove the graph from the set of registered graphs. - registered_graphs_.erase(graph); - - // If no more graphs are registered, deallocate the GPU memory for the seed - // and offset. - if (registered_graphs_.empty()) { - seed_extragraph_.reset(); - offset_extragraph_.reset(); +void CUDAGeneratorState::increase(uint64_t increment) { + // see Note [Why enforce RNG offset % 4 == 0?] + increment = ((increment + 3) / 4) * 4; + + auto capture_id = at::cuda::currentStreamCaptureId(); + if (capture_id.has_value()) { + auto* capture_state = get_capture_state(capture_id.value()); + TORCH_CHECK(capture_state != nullptr, + "RNG op during graph capture but generator is not registered with " + "the capturing graph. Call graph.register_generator_state() before " + "capture_begin()."); + capture_state->increase(increment); + } else { + TORCH_INTERNAL_ASSERT( + philox_offset_per_thread_ % 4 == 0, + "RNG offset must be a multiple of 4."); + philox_offset_per_thread_ += increment; } } -/** - * Note [Explicit Registration of Generators to the CUDA Graph] - * ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - * - * Ideally, it would be more user-friendly if the state could be exchanged and generators - * could be registered with the CUDA graph implicitly. However, resetting GPU tensors during - * the capture stage causes these reset operations to be recorded within the CUDA graph. - * This behavior is undesirable because we do not want these tensors to be reset during - * the replay stage of the graph. - * - * As of now, there is no available method to perform a CUDA operation during the graph's - * recording phase without having that operation be included in the CUDA graph. - * This limitation necessitates explicit user action to register generators with the graph. - * By requiring users to manually register their generators, we can ensure that state resets - * (capture_prologue) only occur before the graph capture begins, thus avoiding unintended - * resets during the replay of the graph. See https://github.com/pytorch/pytorch/pull/114068. - */ - -/** - * Performs the prologue steps for capturing a CUDA graph state. - * This method is intended to reset graph-related state variables before capturing begins. - */ -void CUDAGeneratorState::capture_prologue() { - capturing_ = true; - offset_intragraph_ = 0; - seed_extragraph_.fill_(static_cast(seed_)); - offset_extragraph_.fill_(0); +uint64_t CUDAGeneratorState::capture_epilogue(CaptureId_t capture_id) { + auto* capture_state = get_capture_state(capture_id); + if (capture_state) { + return capture_state->finalize(); + } + return 0; } -/** - * Ends the capturing phase and resets related variables, returning the whole - * graph increment. - */ -uint64_t CUDAGeneratorState::capture_epilogue() { - capturing_ = false; - return offset_intragraph_; +void CUDAGeneratorState::remove_capture_state(CaptureId_t capture_id) { + std::lock_guard lock(capture_states_mutex_); + capture_states_.erase(capture_id); } -/** - * Prepares the state for replay by setting initial state tensors and applying - * total increment. - */ -void CUDAGeneratorState::replay_prologue(uint64_t wholegraph_increment) { - // Ensures the generator is not in capturing mode. - at::cuda::assertNotCapturing( - "Cannot prepare for replay during capturing stage."); - if (wholegraph_increment) { - seed_extragraph_.fill_(static_cast(seed_)); - offset_extragraph_.fill_(static_cast(philox_offset_per_thread_)); - // Applies the total increment achieved during previous captures to update the - // offset. - increase(wholegraph_increment); +void CUDAGeneratorState::replay_prologue(CaptureId_t capture_id, uint64_t wholegraph_increment) { + if (wholegraph_increment == 0) { + return; } + + auto* capture_state = get_capture_state(capture_id); + TORCH_INTERNAL_ASSERT(capture_state != nullptr, + "replay_prologue called but no capture state found for this capture_id"); + capture_state->setup_for_replay(seed_, philox_offset_per_thread_); + philox_offset_per_thread_ += wholegraph_increment; } /** @@ -271,7 +264,8 @@ CUDAGeneratorImpl::CUDAGeneratorImpl( * See Note [Acquire lock when using random generators] */ void CUDAGeneratorImpl::set_current_seed(uint64_t seed) { - if (C10_LIKELY(at::cuda::currentStreamCaptureStatus() == at::cuda::CaptureStatus::None)) { + auto capture_id = at::cuda::currentStreamCaptureId(); + if (C10_LIKELY(!capture_id.has_value())) { state_->seed_ = seed; state_->philox_offset_per_thread_ = 0; no_reset_rnn_state_.clear(); @@ -410,10 +404,14 @@ void CUDAGeneratorImpl::set_philox_offset_per_thread(uint64_t offset) { // set_philox_offset_per_thread instead of set_offset will cause the // cudnn RNN rng state to become stale. TORCH_CHECK(offset % 4 == 0, "offset must be a multiple of 4"); - if (C10_LIKELY(at::cuda::currentStreamCaptureStatus() == at::cuda::CaptureStatus::None)) { + auto capture_id = at::cuda::currentStreamCaptureId(); + if (C10_LIKELY(!capture_id.has_value())) { state_->philox_offset_per_thread_ = offset; } else { - state_->offset_intragraph_ = offset; + auto* capture_state = state_->get_capture_state(capture_id.value()); + TORCH_CHECK(capture_state != nullptr, + "Generator not registered with the capturing graph."); + capture_state->offset_intragraph_ = offset; } } @@ -421,26 +419,19 @@ void CUDAGeneratorImpl::set_philox_offset_per_thread(uint64_t offset) { * Gets the current philox_offset_per_thread_ of CUDAGeneratorImpl. */ uint64_t CUDAGeneratorImpl::philox_offset_per_thread() const { - if (C10_LIKELY(at::cuda::currentStreamCaptureStatus() == at::cuda::CaptureStatus::None)) { + auto capture_id = at::cuda::currentStreamCaptureId(); + if (C10_LIKELY(!capture_id.has_value())) { return state_->philox_offset_per_thread_; } else { - return state_->offset_intragraph_; + auto* capture_state = state_->get_capture_state(capture_id.value()); + TORCH_CHECK(capture_state != nullptr, + "Generator not registered with the capturing graph."); + return capture_state->offset_intragraph_; } } -/** - * Registers this state to a CUDA graph to manage within the graph. - */ void CUDAGeneratorImpl::register_graph(cuda::CUDAGraph* graph) { graph->register_generator_state(state_); - state_->register_graph(graph); -} - -/** - * Unregisters a CUDA graph from the RNG state. - */ -void CUDAGeneratorImpl::unregister_graph(cuda::CUDAGraph* graph) { - state_->unregister_graph(graph); } /** @@ -465,12 +456,20 @@ void CUDAGeneratorImpl::unregister_graph(cuda::CUDAGraph* graph) { * See Note [Acquire lock when using random generators] */ PhiloxCudaState CUDAGeneratorImpl::philox_cuda_state(uint64_t increment) { - if (at::cuda::currentStreamCaptureStatus() != at::cuda::CaptureStatus::None) { - uint64_t offset = state_->offset_intragraph_; + auto capture_id = at::cuda::currentStreamCaptureId(); + if (capture_id.has_value()) { + auto* capture_state = state_->get_capture_state(capture_id.value()); + TORCH_CHECK(capture_state != nullptr, + "RNG op during graph capture but generator is not registered with " + "the capturing graph. Call graph.register_generator_state() before " + "capture_begin()."); + + uint64_t offset = capture_state->offset_intragraph_; state_->increase(increment); + return PhiloxCudaState( - state_->seed_extragraph_.data_ptr(), - state_->offset_extragraph_.data_ptr(), + capture_state->rng_state_seed_extragraph_.data_ptr(), + capture_state->rng_state_offset_extragraph_.data_ptr(), offset); } else { uint64_t offset = state_->philox_offset_per_thread_; diff --git a/aten/src/ATen/cuda/CUDAGeneratorImpl.h b/aten/src/ATen/cuda/CUDAGeneratorImpl.h index d4ab49382e7ff..fe62fc53eef07 100644 --- a/aten/src/ATen/cuda/CUDAGeneratorImpl.h +++ b/aten/src/ATen/cuda/CUDAGeneratorImpl.h @@ -4,15 +4,20 @@ #include #include #include +#include +#include #include #include -#include +#include + namespace at { namespace cuda { struct CUDAGraph; } +using CaptureId_t = c10::CaptureId_t; + /** * Note [CUDA Graph-safe RNG states] * ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -93,32 +98,51 @@ struct CUDAGraph; * */ +/** + * Per-capture state for a generator. + * Each (generator, capture_id) pair gets its own CUDAGeneratorCaptureState. + * This holds the GPU tensors and offset tracking for a specific graph capture. + */ +struct CUDAGeneratorCaptureState : public c10::intrusive_ptr_target { + uint64_t offset_intragraph_{0}; + at::TensorBase rng_state_seed_extragraph_; + at::TensorBase rng_state_offset_extragraph_; + + CUDAGeneratorCaptureState() = default; + + bool is_initialized() const { return rng_state_seed_extragraph_.defined(); } + void initialize(uint64_t seed); + void increase(uint64_t increment); + uint64_t finalize(); + void setup_for_replay(uint64_t seed, uint64_t philox_offset); +}; + +/** + * Generator state that supports multiple concurrent graph captures. + * Each capture gets its own CUDAGeneratorCaptureState keyed by CaptureId_t. + */ struct CUDAGeneratorState : public c10::intrusive_ptr_target { uint64_t seed_; uint64_t philox_offset_per_thread_; - uint64_t offset_intragraph_; - bool capturing_{}; - std::unordered_set registered_graphs_; - at::TensorBase seed_extragraph_; - at::TensorBase offset_extragraph_; + + // Map from capture ID to per-capture state + ska::flat_hash_map> capture_states_; + mutable std::mutex capture_states_mutex_; CUDAGeneratorState( uint64_t seed = default_rng_seed_val, - uint64_t philox_offset_per_thread = 0, - uint64_t offset_intragraph = 0) + uint64_t philox_offset_per_thread = 0) : seed_(seed), - philox_offset_per_thread_(philox_offset_per_thread), - offset_intragraph_(offset_intragraph) {} + philox_offset_per_thread_(philox_offset_per_thread) {} void increase(uint64_t increment); - void register_graph(cuda::CUDAGraph* graph); - void unregister_graph(cuda::CUDAGraph* graph); + CUDAGeneratorCaptureState* get_capture_state(CaptureId_t capture_id); + void init_capture_state(CaptureId_t capture_id); + uint64_t capture_epilogue(CaptureId_t capture_id); + void replay_prologue(CaptureId_t capture_id, uint64_t wholegraph_increment); + void remove_capture_state(CaptureId_t capture_id); - void capture_prologue(); - // capture_epilogue returns the wholegraph_increment - uint64_t capture_epilogue(); - void replay_prologue(uint64_t wholegraph_increment); c10::intrusive_ptr clone(); }; @@ -147,7 +171,6 @@ struct TORCH_CUDA_CPP_API CUDAGeneratorImpl : public c10::GeneratorImpl { uint64_t philox_offset_per_thread() const; void register_graph(cuda::CUDAGraph* graph); - void unregister_graph(cuda::CUDAGraph* graph); // Generates a PhiloxCudaState with a specified increment, and increment // current state diff --git a/aten/src/ATen/cuda/CUDAGraph.cpp b/aten/src/ATen/cuda/CUDAGraph.cpp index f1e25cfe011a6..b58e60887b299 100644 --- a/aten/src/ATen/cuda/CUDAGraph.cpp +++ b/aten/src/ATen/cuda/CUDAGraph.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include #include #include @@ -9,6 +10,7 @@ #include #include +#include namespace at::cuda { @@ -78,6 +80,24 @@ void CUDAGraph::register_generator_state(const at::Generator& generator) { cuda_gen->register_graph(this); } + +template <> +std::function CUDAGraph::create_allocate_filter() const { + return [this](cudaStream_t stream) { + auto capture_id_opt = c10::cuda::captureIdMayInitCtx(stream); + return capture_id_opt.has_value() && capture_id_opt.value() == capture_id_; + }; +} + +template <> +std::function CUDAGraph::create_allocate_filter() const { + return [this](c10::Stream stream) { + cudaStream_t cuda_stream = CUDAStream(CUDAStream::UNCHECKED, stream); + auto capture_id_opt = c10::cuda::captureIdMayInitCtx(cuda_stream); + return capture_id_opt.has_value() && capture_id_opt.value() == capture_id_; + }; +} + void CUDAGraph::capture_begin(MempoolId_t pool/*={0,0}*/, cudaStreamCaptureMode capture_mode) { TORCH_CHECK(!has_graph_exec_, "This CUDAGraph instance already owns a captured graph. " @@ -90,11 +110,6 @@ void CUDAGraph::capture_begin(MempoolId_t pool/*={0,0}*/, cudaStreamCaptureMode std::nullopt, cuda::detail::getDefaultCUDAGenerator()); gen->register_graph(this); - for (auto& [generator_state, wholegraph_increments] : - captured_generator_states_) { - generator_state->capture_prologue(); - } - auto stream = at::cuda::getCurrentCUDAStream(); TORCH_CHECK(stream != at::cuda::getDefaultCUDAStream(), @@ -122,27 +137,29 @@ void CUDAGraph::capture_begin(MempoolId_t pool/*={0,0}*/, cudaStreamCaptureMode // Addendum: beginAllocateStreamToPool is now called before cudaStreamBeginCapture to prevent an // autograd thread's free() call triggering an invalid cudaEventRecord in the caching allocator // due to the capture status being updated _after_ a capture had already started. - c10::cuda::CUDACachingAllocator::beginAllocateToPool(capture_dev_, mempool_id_, create_allocate_filter()); + c10::cuda::CUDACachingAllocator::beginAllocateToPool(capture_dev_, mempool_id_, create_allocate_filter()); - auto filter = create_allocate_filter(); - - at::getHostAllocator(at::kCUDA)->begin_allocate_to_pool(mempool_id_, [filter](c10::Stream stream) { - return filter(CUDAStream(CUDAStream::UNCHECKED, stream)); - }); + at::getHostAllocator(at::kCUDA)->begin_allocate_to_pool(mempool_id_, create_allocate_filter()); // cudaStreamCaptureModeGlobal is the most conservative option to // prevent potentially unsafe CUDA API calls during capture. See // https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__STREAM.html#group__CUDART__STREAM_1g9d0535d93a214cbf126835257b16ba85 AT_CUDA_CHECK(cudaStreamBeginCapture(capture_stream_, capture_mode)); - cudaStreamCaptureStatus status{}; - AT_CUDA_CHECK(cudaStreamGetCaptureInfo(stream, &status, &capture_id_)); - TORCH_INTERNAL_ASSERT(status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive); + auto capture_id_opt = c10::cuda::captureIdMayInitCtx(stream); + TORCH_INTERNAL_ASSERT(capture_id_opt.has_value(), + "Stream should be actively capturing after cudaStreamBeginCapture"); + capture_id_ = capture_id_opt.value(); { - std::unique_lock lock(_currently_capturing_graphs_mutex); + std::lock_guard lock(_currently_capturing_graphs_mutex); _currently_capturing_graphs.emplace(capture_id_, this); } + + for (auto& [generator_state, wholegraph_increment] : + captured_generator_states_) { + generator_state->init_capture_state(capture_id_); + } } void CUDAGraph::capture_end() { @@ -169,9 +186,9 @@ void CUDAGraph::capture_end() { TORCH_CHECK(graph_ != nullptr, "Invalid capture."); - for (auto& [generator_state, wholegraph_increments] : + for (auto& [generator_state, wholegraph_increment] : captured_generator_states_) { - wholegraph_increments = generator_state->capture_epilogue(); + wholegraph_increment = generator_state->capture_epilogue(capture_id_); } size_t numCUDAGraphNodes = 0; @@ -233,9 +250,9 @@ void CUDAGraph::replay() { c10::OptionalDeviceGuard device_guard{capture_stream_.device()}; - for (auto& [generator_state, wholegraph_increments] : + for (auto& [generator_state, wholegraph_increment] : captured_generator_states_) { - generator_state->replay_prologue(wholegraph_increments); + generator_state->replay_prologue(capture_id_, wholegraph_increment); } // graph_exec_ may be replayed in any stream. AT_CUDA_CHECK(cudaGraphLaunch(graph_exec_, at::cuda::getCurrentCUDAStream())); @@ -294,6 +311,20 @@ void CUDAGraph::reset() { // and the allocator could end up in all kinds of weird states depending where failure occurred. // If the user catches the failure exception in a script, or is running in REPL or (god forbid) // a Jupyter notebook, I don't see an easy way for reset() to gracefully fix all such possible error states. + + if (capture_id_ != 0) { + for (auto& [generator_state, wholegraph_increment] : captured_generator_states_) { + generator_state->remove_capture_state(capture_id_); + } + } + captured_generator_states_.clear(); + + if (capture_id_ != 0) { + std::lock_guard lock(_currently_capturing_graphs_mutex); + _currently_capturing_graphs.erase(capture_id_); + capture_id_ = 0; + } + if (capture_ended_) { // Clean up cuBLAS workspaces allocated on the capture stream, otherwise live allocations prevent // private pool cleanup @@ -322,10 +353,6 @@ MempoolId_t CUDAGraph::pool() { } CUDAGraph::~CUDAGraph() { - for (auto& [generator_state, wholegraph_increments] : - captured_generator_states_) { - generator_state->unregister_graph(this); - } reset(); // There are recent HIP changes where hipGraphExecDestroy doesn't immediately free memory. @@ -343,17 +370,14 @@ CUDAGraph::~CUDAGraph() { CUDAGraph* CUDAGraph::get_currently_capturing_graph() { std::unique_lock lock(_currently_capturing_graphs_mutex); - cudaStreamCaptureStatus status{}; - CaptureId_t current_capture_id = 0; - auto stream = at::cuda::getCurrentCUDAStream(); - AT_CUDA_CHECK(cudaStreamGetCaptureInfo(stream, &status, ¤t_capture_id)); + auto capture_id_opt = c10::cuda::currentStreamCaptureIdMayInitCtx(); TORCH_CHECK( - status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive, + capture_id_opt.has_value(), "The current stream is not currently capturing."); TORCH_CHECK( - _currently_capturing_graphs.count(current_capture_id), + _currently_capturing_graphs.count(capture_id_opt.value()), "get_currently_capturing_graph() can be used only between capture_begin() and capture_end(). Did you use a stream without making it depend upon the original stream used for capture?"); - return _currently_capturing_graphs.at(current_capture_id); + return _currently_capturing_graphs.at(capture_id_opt.value()); } void CUDAGraph::begin_capture_to_if_node( @@ -439,13 +463,6 @@ getCurrentCUDAStream(), &cond_node, nullptr, 1, cudaStreamSetCaptureDependencies CUDAStream child_stream = getStreamFromPool(); conditional_graph_capture_ids_.push(0); - conditional_rng_snapshots_.emplace(); - auto& conditional_rng_snapshot = conditional_rng_snapshots_.top(); - for (auto& [generator_state, wholegraph_increments] : - captured_generator_states_) { - conditional_rng_snapshot.emplace( - generator_state, generator_state->offset_intragraph_); - } c10::cuda::CUDACachingAllocator::endAllocateToPool(capture_dev_, mempool_id_); at::getHostAllocator(at::kCUDA)->end_allocate_to_pool(mempool_id_); @@ -459,10 +476,10 @@ getCurrentCUDAStream(), &cond_node, nullptr, 1, cudaStreamSetCaptureDependencies AT_CUDA_CHECK(cudaStreamBeginCaptureToGraph( child_stream, if_node_child_graph, nullptr, nullptr, 0, capture_mode_)); - AT_CUDA_CHECK(cudaStreamGetCaptureInfo( - child_stream, &status, &conditional_graph_capture_ids_.top())); - TORCH_INTERNAL_ASSERT( - status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive); + auto child_capture_id_opt = c10::cuda::captureIdMayInitCtx(child_stream); + TORCH_INTERNAL_ASSERT(child_capture_id_opt.has_value(), + "Child stream should be actively capturing after cudaStreamBeginCaptureToGraph"); + conditional_graph_capture_ids_.top() = child_capture_id_opt.value(); conditional_node_streams_.emplace(child_stream); @@ -483,50 +500,38 @@ getCurrentCUDAStream(), &cond_node, nullptr, 1, cudaStreamSetCaptureDependencies void CUDAGraph::end_capture_to_conditional_node() { #if !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040) TORCH_INTERNAL_ASSERT( - !conditional_rng_snapshots_.empty(), - "Missing RNG snapshot for conditional node capture."); + !conditional_graph_capture_ids_.empty(), + "Missing capture ID for conditional node."); + CaptureId_t child_capture_id = conditional_graph_capture_ids_.top(); bool rng_or_generators_changed = false; - auto& conditional_rng_snapshot = conditional_rng_snapshots_.top(); - if (conditional_rng_snapshot.size() != captured_generator_states_.size()) { - rng_or_generators_changed = true; - } else { - for (const auto& [generator_state, offset_intragraph_before_capture] : - conditional_rng_snapshot) { - const auto generator_it = captured_generator_states_.find(generator_state); - if (generator_it == captured_generator_states_.end() || - generator_state->offset_intragraph_ != - offset_intragraph_before_capture) { - rng_or_generators_changed = true; - break; - } + for (const auto& [generator_state, wholegraph_increment] : + captured_generator_states_) { + if (generator_state->get_capture_state(child_capture_id) != nullptr) { + rng_or_generators_changed = true; + break; } } { std::unique_lock lock(_currently_capturing_graphs_mutex); - CaptureId_t capture_id = conditional_graph_capture_ids_.top(); TORCH_CHECK( - _currently_capturing_graphs.count(capture_id), + _currently_capturing_graphs.count(child_capture_id), "capture_end() called before capture_begin()."); - _currently_capturing_graphs.erase(capture_id); + _currently_capturing_graphs.erase(child_capture_id); } CUDAStream stream = conditional_node_streams_.top().current_stream(); AT_CUDA_CHECK(cudaStreamEndCapture(stream.stream(), nullptr)); conditional_node_streams_.pop(); conditional_graph_capture_ids_.pop(); - conditional_rng_snapshots_.pop(); c10::cuda::CUDACachingAllocator::endAllocateToPool(capture_dev_, mempool_id_); at::getHostAllocator(at::kCUDA)->end_allocate_to_pool(mempool_id_); if (conditional_graph_capture_ids_.empty()) { c10::cuda::CUDACachingAllocator::beginAllocateToPool( - capture_dev_, mempool_id_, create_allocate_filter()); - auto filter = create_allocate_filter(); - at::getHostAllocator(at::kCUDA)->begin_allocate_to_pool(mempool_id_, [filter](c10::Stream stream) { - return filter(CUDAStream(CUDAStream::UNCHECKED, stream)); - }); + capture_dev_, mempool_id_, create_allocate_filter()); + at::getHostAllocator(at::kCUDA)->begin_allocate_to_pool(mempool_id_, create_allocate_filter()); } else { c10::cuda::CUDACachingAllocator::beginAllocateToPool( capture_dev_, mempool_id_, create_child_allocate_filter()); @@ -535,7 +540,6 @@ void CUDAGraph::end_capture_to_conditional_node() { return filter(CUDAStream(CUDAStream::UNCHECKED, stream)); }); } - constexpr const char* rng_with_conditional_nodes_error = "RNG within data-dependent conditional nodes is not supported yet."; TORCH_CHECK(!rng_or_generators_changed, rng_with_conditional_nodes_error); @@ -547,22 +551,11 @@ void CUDAGraph::end_capture_to_conditional_node() { #endif } -std::function CUDAGraph::create_allocate_filter() { - return [this](cudaStream_t stream) { - cudaStreamCaptureStatus status{}; - CaptureId_t stream_capture_id = 0; - AT_CUDA_CHECK(cudaStreamGetCaptureInfo(stream, &status, &stream_capture_id)); - return status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive && stream_capture_id == capture_id_; - }; -} - std::function CUDAGraph::create_child_allocate_filter() { #if !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040) return [¤t_capture_id = conditional_graph_capture_ids_.top()](cudaStream_t stream) { - cudaStreamCaptureStatus status{}; - CaptureId_t stream_capture_id{}; - AT_CUDA_CHECK(cudaStreamGetCaptureInfo(stream, &status, &stream_capture_id)); - return status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive && stream_capture_id == current_capture_id; + auto capture_id_opt = c10::cuda::captureIdMayInitCtx(stream); + return capture_id_opt.has_value() && capture_id_opt.value() == current_capture_id; }; #else // !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040) AT_ERROR( diff --git a/aten/src/ATen/cuda/CUDAGraph.h b/aten/src/ATen/cuda/CUDAGraph.h index edaad31b42bb5..9a971e82bb787 100644 --- a/aten/src/ATen/cuda/CUDAGraph.h +++ b/aten/src/ATen/cuda/CUDAGraph.h @@ -9,6 +9,7 @@ #include #include +#include #include #if defined(USE_ROCM) || !(defined(CUDA_VERSION) && CUDA_VERSION >= 12040) @@ -64,7 +65,6 @@ struct TORCH_CUDA_CPP_API CUDAGraph { CUDAGraph(CUDAGraph&& other) = delete; CUDAGraph& operator=(CUDAGraph&& other) = delete; - // See Note [Explicit Registration of Generators to the CUDA Graph] void register_generator_state(c10::intrusive_ptr state); void register_generator_state(const at::Generator& generator); void capture_begin( @@ -88,7 +88,8 @@ struct TORCH_CUDA_CPP_API CUDAGraph { const Tensor& scalar_cuda_pred_tensor); private: - std::function create_allocate_filter(); + template + std::function create_allocate_filter() const; std::function create_child_allocate_filter(); protected: @@ -146,11 +147,13 @@ struct TORCH_CUDA_CPP_API CUDAGraph { #if !defined(USE_ROCM) && (defined(CUDA_VERSION) && CUDA_VERSION >= 12040) std::stack conditional_node_streams_; std::stack conditional_graph_capture_ids_; - std::stack< - ska::flat_hash_map, uint64_t>> - conditional_rng_snapshots_; #endif // !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040 }; +template <> +std::function CUDAGraph::create_allocate_filter() const; +template <> +std::function CUDAGraph::create_allocate_filter() const; + } // namespace cuda } // namespace at diff --git a/aten/src/ATen/cuda/CUDAGraphsUtils.cuh b/aten/src/ATen/cuda/CUDAGraphsUtils.cuh index d3a5b306eeea4..64ab0140793bb 100644 --- a/aten/src/ATen/cuda/CUDAGraphsUtils.cuh +++ b/aten/src/ATen/cuda/CUDAGraphsUtils.cuh @@ -19,12 +19,17 @@ using CaptureStatus = c10::cuda::CaptureStatus; // Use this version where you don't want to create a CUDA context if none exists. inline CaptureStatus currentStreamCaptureStatus() { - // don't create a context if we don't have to if (c10::cuda::hasPrimaryContext(c10::cuda::current_device())) { return c10::cuda::currentStreamCaptureStatusMayInitCtx(); - } else { - return CaptureStatus::None; } + return CaptureStatus::None; +} + +inline std::optional currentStreamCaptureId() { + if (c10::cuda::hasPrimaryContext(c10::cuda::current_device())) { + return c10::cuda::currentStreamCaptureIdMayInitCtx(); + } + return std::nullopt; } inline void assertNotCapturing(const std::string& attempt) { diff --git a/aten/src/ATen/cuda/CUDAGreenContext.cpp b/aten/src/ATen/cuda/CUDAGreenContext.cpp index 7a0eb81bfe3c2..07861d7d78a11 100644 --- a/aten/src/ATen/cuda/CUDAGreenContext.cpp +++ b/aten/src/ATen/cuda/CUDAGreenContext.cpp @@ -11,10 +11,27 @@ C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wunused-private-field") #endif +#if defined(CUDA_VERSION) && CUDA_VERSION >= 13010 && HAS_CUDA_GREEN_CONTEXT() +#define HAS_CUDA_WORKQUEUE_SUPPORT() 1 +#else +#define HAS_CUDA_WORKQUEUE_SUPPORT() 0 +#endif + namespace at::cuda { -GreenContext::GreenContext(uint32_t device_id, uint32_t num_sms) { +GreenContext::GreenContext( + uint32_t device_id, + std::optional num_sms, + std::optional workqueue_scope, + std::optional workqueue_concurrency_limit) { #if HAS_CUDA_GREEN_CONTEXT() + TORCH_CHECK( + num_sms.has_value() || workqueue_scope.has_value(), + "At least one of num_sms or workqueue_scope must be specified"); + TORCH_CHECK( + !workqueue_concurrency_limit.has_value() || workqueue_scope.has_value(), + "workqueue_concurrency_limit requires workqueue_scope to be set"); + int driver_version; C10_CUDA_CHECK(cudaDriverGetVersion(&driver_version)); TORCH_CHECK( @@ -29,46 +46,73 @@ GreenContext::GreenContext(uint32_t device_id, uint32_t num_sms) { cudaFree(nullptr); } - CUdevice device; + CUdevice device; device_id_ = device_id; C10_CUDA_DRIVER_CHECK( c10::cuda::DriverAPI::get()->cuDeviceGet_(&device, device_id)); - // Get device resources - CUdevResource device_resource; - C10_CUDA_DRIVER_CHECK(c10::cuda::DriverAPI::get()->cuDeviceGetDevResource_( - device, &device_resource, CU_DEV_RESOURCE_TYPE_SM)); + std::vector resources; - TORCH_CHECK( - num_sms > 0 && num_sms <= device_resource.sm.smCount, - "Invalid number of SMs requested for green context: ", - num_sms, - " (device has ", - device_resource.sm.smCount, - " SMs)"); - - // Split resources - std::vector result(1); - auto result_data = result.data(); - unsigned int nb_groups = 1; - CUdevResource remaining; + // --- SM resource --- + if (num_sms.has_value()) { + CUdevResource sm_resource; + C10_CUDA_DRIVER_CHECK(c10::cuda::DriverAPI::get()->cuDeviceGetDevResource_( + device, &sm_resource, CU_DEV_RESOURCE_TYPE_SM)); - C10_CUDA_DRIVER_CHECK( - c10::cuda::DriverAPI::get()->cuDevSmResourceSplitByCount_( - result_data, - &nb_groups, - &device_resource, - &remaining, - 0, // default flags - num_sms)); + TORCH_CHECK( + *num_sms > 0 && *num_sms <= sm_resource.sm.smCount, + "Invalid number of SMs requested for green context: ", + *num_sms, + " (device has ", + sm_resource.sm.smCount, + " SMs)"); + + // Split resources + std::vector split_result(1); + unsigned int nb_groups = 1; + CUdevResource remaining; + + C10_CUDA_DRIVER_CHECK( + c10::cuda::DriverAPI::get()->cuDevSmResourceSplitByCount_( + split_result.data(), + &nb_groups, + &sm_resource, + &remaining, + 0, // default flags + *num_sms)); + TORCH_CHECK(nb_groups == 1, "Failed to create single SM resource group"); + resources.push_back(split_result[0]); + } - TORCH_CHECK(nb_groups == 1, "Failed to create single resource group"); + // --- Workqueue config resource --- + if (workqueue_scope.has_value()) { +#if HAS_CUDA_WORKQUEUE_SUPPORT() + TORCH_CHECK( + driver_version >= 13010, "cuda driver too old to use workqueue configuration!"); + CUdevResource wq_resource{}; + C10_CUDA_DRIVER_CHECK(c10::cuda::DriverAPI::get()->cuDeviceGetDevResource_( + device, &wq_resource, CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG)); + + wq_resource.wqConfig.sharingScope = + static_cast(*workqueue_scope); + if (workqueue_concurrency_limit.has_value()) { + wq_resource.wqConfig.wqConcurrencyLimit = *workqueue_concurrency_limit; + } + resources.push_back(wq_resource); +#else + TORCH_CHECK( + false, + "Workqueue configuration for green contexts requires CUDA 13.1+!"); +#endif + } // Generate resource descriptor CUdevResourceDesc desc; C10_CUDA_DRIVER_CHECK( c10::cuda::DriverAPI::get()->cuDevResourceGenerateDesc_( - &desc, result_data, 1)); + &desc, + resources.data(), + static_cast(resources.size()))); // Create green context // CU_GREEN_CTX_DEFAULT_STREAM is required per docs: @@ -83,20 +127,45 @@ GreenContext::GreenContext(uint32_t device_id, uint32_t num_sms) { #else TORCH_CHECK(false, "Green Context is only supported on CUDA 12.8+!"); #endif - } +} - std::unique_ptr GreenContext::create( - uint32_t num_sms, - std::optional device_id) { +std::unique_ptr GreenContext::create( + std::optional device_id, + std::optional num_sms, + std::optional workqueue_scope, + std::optional workqueue_concurrency_limit) { #if HAS_CUDA_GREEN_CONTEXT() - if (!device_id.has_value()) { - device_id = at::cuda::current_device(); - } - return std::unique_ptr(new GreenContext(device_id.value(), num_sms)); + if (!device_id.has_value()) { + device_id = at::cuda::current_device(); + } + return std::unique_ptr(new GreenContext( + device_id.value(), num_sms, workqueue_scope, workqueue_concurrency_limit)); #else - TORCH_CHECK(false, "Green Context is only supported on CUDA 12.8+!"); + TORCH_CHECK(false, "Green Context is only supported on CUDA 12.8+!"); #endif +} + +uint32_t GreenContext::max_workqueue_concurrency( + std::optional device_id) { +#if HAS_CUDA_WORKQUEUE_SUPPORT() + int driver_version; + C10_CUDA_CHECK(cudaDriverGetVersion(&driver_version)); + TORCH_CHECK( + driver_version >= 13010, "cuda driver too old to use workqueue configuration!"); + if (!device_id.has_value()) { + device_id = at::cuda::current_device(); } + CUdevice device; + C10_CUDA_DRIVER_CHECK( + c10::cuda::DriverAPI::get()->cuDeviceGet_(&device, device_id.value())); + CUdevResource wq_resource; + C10_CUDA_DRIVER_CHECK(c10::cuda::DriverAPI::get()->cuDeviceGetDevResource_( + device, &wq_resource, CU_DEV_RESOURCE_TYPE_WORKQUEUE_CONFIG)); + return wq_resource.wqConfig.wqConcurrencyLimit; +#else + TORCH_CHECK(false, "Workqueue configuration requires CUDA 13.1+!"); +#endif +} // Implement move operations #if HAS_CUDA_GREEN_CONTEXT() diff --git a/aten/src/ATen/cuda/CUDAGreenContext.h b/aten/src/ATen/cuda/CUDAGreenContext.h index b1eda3f48a45a..a93eedfaeedfa 100644 --- a/aten/src/ATen/cuda/CUDAGreenContext.h +++ b/aten/src/ATen/cuda/CUDAGreenContext.h @@ -11,12 +11,24 @@ namespace { constexpr int kStreamPerGreenContextPool = 32; } +// Workqueue sharing scope for green contexts. +// Values match the CUDA driver API's CUdevWorkqueueConfigScope enum. +enum class WorkqueueScope : int32_t { + DeviceCtx = 0, + Balanced = 1, +}; + class TORCH_CUDA_CPP_API GreenContext { public: - // Green context creation static std::unique_ptr create( - uint32_t num_sms, - std::optional device_id); + std::optional device_id, + std::optional num_sms, + std::optional workqueue_scope = std::nullopt, + std::optional workqueue_concurrency_limit = std::nullopt); + + static uint32_t max_workqueue_concurrency( + std::optional device_id = std::nullopt); + ~GreenContext() noexcept; // Delete copy constructor and assignment @@ -31,7 +43,12 @@ class TORCH_CUDA_CPP_API GreenContext { CUDAStream Stream(); private: - GreenContext(uint32_t device_id, uint32_t num_sms); + GreenContext( + uint32_t device_id, + std::optional num_sms, + std::optional workqueue_scope, + std::optional workqueue_concurrency_limit); + // Implement move operations GreenContext(GreenContext&& other) noexcept; GreenContext& operator=(GreenContext&& other) noexcept; diff --git a/aten/src/ATen/cuda/CachingHostAllocator.cpp b/aten/src/ATen/cuda/CachingHostAllocator.cpp index dcff81333cb2f..c376eb3a0a4d5 100644 --- a/aten/src/ATen/cuda/CachingHostAllocator.cpp +++ b/aten/src/ATen/cuda/CachingHostAllocator.cpp @@ -3,6 +3,8 @@ #include #include #include +#include +#include #include #include @@ -14,6 +16,25 @@ using Block = HostBlock; struct CUDACachingHostAllocatorImpl : public CachingHostAllocatorImpl { + void free(void* ctx) override { + using Base = CachingHostAllocatorImpl; + try { + Base::free(ctx); + } catch (...) { + if (!c10::cuda::CUDACachingAllocator::CUDAAllocatorConfig:: + pinned_free_catch_all()) { + TORCH_WARN("Exception in pinned allocator free(), rethrowing"); + throw; + } + // pinned_free_catch_all is enabled: suppress the exception to prevent + // it from escaping through ~StorageImpl() (implicitly noexcept), which + // would cause std::terminate. Allows graceful shutdown to proceed. + STATIC_GAUGE(pytorch.CUDACachingHostAllocator.free_fail_catch_all) + .record(1); + TORCH_WARN("Suppressed exception in pinned allocator free()"); + } + } + private: ska::flat_hash_map use_host_register; @@ -245,9 +266,7 @@ struct CUDACachingHostAllocatorImpl } bool stream_is_capturing(CUDAStream s) const override { - cudaStreamCaptureStatus status{cudaStreamCaptureStatusNone}; - C10_CUDA_CHECK(cudaStreamIsCapturing(s, &status)); - return status != cudaStreamCaptureStatusNone; + return s.is_capturing(); } }; diff --git a/aten/src/ATen/cuda/CublasHandlePool.cpp b/aten/src/ATen/cuda/CublasHandlePool.cpp index f3df32749fa6b..1fbe3cdedf474 100644 --- a/aten/src/ATen/cuda/CublasHandlePool.cpp +++ b/aten/src/ATen/cuda/CublasHandlePool.cpp @@ -7,9 +7,14 @@ #include #include #include +#include #include #include +#if defined(USE_ROCM) +#include +#endif + /** * Note [hipblaslt handles] * ~~~~~~~~~~~~~~~~~~~~~~~~ @@ -24,6 +29,10 @@ * For CUDA builds, getCurrentCUDABlasLtHandle will alias for getCurrentCUDABlasHandle, * whereas for ROCm builds, it is a distinct function. * + * Additionally, hipblaslt cannot share a single handle across multiple streams. + * On ROCm, getCurrentCUDABlasLtHandle returns a handle unique to each (device, stream) + * pair, rather than just per-device like the cublas handle pool. + * * The workspace pools are separate for ROCm. On CUDA, the env var * TORCH_CUBLASLT_UNIFIED_WORKSPACE can be used to opt-in to unifying the workspace pools. */ @@ -53,8 +62,6 @@ void destroyCublasLtHandle(cublasLtHandle_t handle) { using CuBlasLtPoolType = DeviceThreadHandlePool; // ugly hack until hipblasSetWorkspace exists -#include - static hipblasStatus_t rocBLASStatusToHIPStatus(rocblas_status error) { switch(error) { case rocblas_status_size_unchanged: @@ -135,24 +142,16 @@ void clearCublasWorkspacesForStream(cudaStream_t stream) { { auto& workspace = cublas_handle_stream_to_workspace(); std::unique_lock lock(workspace.mutex); - for (auto it = workspace.map.begin(); it != workspace.map.end(); ) { - if (std::get<1>(it->first) == stream_ptr) { - it = workspace.map.erase(it); - } else { - ++it; - } - } + std::erase_if(workspace.map, [stream_ptr](const auto& entry) { + return std::get<1>(entry.first) == stream_ptr; + }); } { auto& workspace = cublaslt_handle_stream_to_workspace(); std::unique_lock lock(workspace.mutex); - for (auto it = workspace.map.begin(); it != workspace.map.end(); ) { - if (std::get<1>(it->first) == stream_ptr) { - it = workspace.map.erase(it); - } else { - ++it; - } - } + std::erase_if(workspace.map, [stream_ptr](const auto& entry) { + return std::get<1>(entry.first) == stream_ptr; + }); } } @@ -248,10 +247,10 @@ size_t parseCUDABlasLtWorkspaceSize() { } size_t getCUDABlasLtWorkspaceSize() { - size_t pool_size = parseCUDABlasLtWorkspaceSize(); + static size_t pool_size = parseCUDABlasLtWorkspaceSize(); #ifndef USE_ROCM if (unified_cublas_and_lt_workspaces()) { - auto cublasWorkspaceSize = parseChosenWorkspaceSize(); + static size_t cublasWorkspaceSize = parseChosenWorkspaceSize(); if (cublasWorkspaceSize < pool_size) { TORCH_WARN_ONCE("Requested unified CUBLASLT workspace size of ", pool_size, " bytes exceeds CUBLAS workspace size of ", cublasWorkspaceSize, @@ -309,15 +308,26 @@ void setWorkspaceForHandle(cublasHandle_t handle, c10::cuda::CUDAStream stream) void* getCUDABlasLtWorkspace() { #ifndef USE_ROCM if (unified_cublas_and_lt_workspaces()) { - cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle(); + cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle(/*setup=*/false); auto stream = c10::cuda::getCurrentCUDAStream(); cudaStream_t _stream = stream; auto key = std::make_tuple(static_cast(handle), static_cast(_stream)); auto& workspace = at::cuda::cublas_handle_stream_to_workspace(); - std::shared_lock lock(workspace.mutex); - auto workspace_it = workspace.map.find(key); - TORCH_INTERNAL_ASSERT(workspace_it != workspace.map.end()); - return workspace_it->second.mutable_get(); + { + std::shared_lock lock(workspace.mutex); + auto workspace_it = workspace.map.find(key); + if (workspace_it != workspace.map.end()) { + return workspace_it->second.mutable_get(); + } + } + // First use for this handle+stream pair — allocate and insert directly. + // No need to call cublasSetWorkspace; Lt passes workspace explicitly. + auto new_workspace = getNewWorkspace(); + { + std::unique_lock lock(workspace.mutex); + auto workspace_it = workspace.map.try_emplace(key, std::move(new_workspace)).first; + return workspace_it->second.mutable_get(); + } } #endif cublasLtHandle_t handle = getCurrentCUDABlasLtHandle(); @@ -349,7 +359,7 @@ void* getCUDABlasLtWorkspace() { } } -cublasHandle_t getCurrentCUDABlasHandle() { +cublasHandle_t getCurrentCUDABlasHandle(bool setup) { c10::DeviceIndex device = 0; AT_CUDA_CHECK(c10::cuda::GetDevice(&device)); @@ -357,8 +367,6 @@ cublasHandle_t getCurrentCUDABlasHandle() { CUcontext pctx = nullptr; at::globalContext().getNVRTC().cuCtxGetCurrent(&pctx); if (C10_UNLIKELY(!pctx)) { - // workaround for corner case where a primary context exists but is not - // the current context, seen in multithreaded use-cases TORCH_WARN_ONCE("Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context..."); at::globalContext().getNVRTC().cuDevicePrimaryCtxRetain(&pctx, device); at::globalContext().getNVRTC().cuCtxSetCurrent(pctx); @@ -381,6 +389,11 @@ cublasHandle_t getCurrentCUDABlasHandle() { pool->newPoolWindow()); auto handle = myPoolWindow->reserve(device); + + if (!setup) { + return handle; + } + auto stream = c10::cuda::getCurrentCUDAStream(); TORCH_CUDABLAS_CHECK(cublasSetStream(handle, stream)); // We explicitly set the cublas workspace even though CUDA 12.2+ fixed the @@ -434,10 +447,14 @@ cublasLtHandle_t getCurrentCUDABlasLtHandle() { thread_local std::unique_ptr myPoolWindow( pool->newPoolWindow()); - auto handle = myPoolWindow->reserve(device); + // hipblaslt cannot share a single handle across multiple streams, + // so reserve a handle unique to each (device, stream) pair. + auto stream = c10::cuda::getCurrentCUDAStream(); + cudaStream_t _stream = stream; + auto handle = myPoolWindow->reserve(device, static_cast(_stream)); return handle; #else - return reinterpret_cast(getCurrentCUDABlasHandle()); + return reinterpret_cast(getCurrentCUDABlasHandle(/*setup=*/false)); #endif } diff --git a/aten/src/ATen/cuda/StatelessPhilox4x32.cuh b/aten/src/ATen/cuda/StatelessPhilox4x32.cuh new file mode 100644 index 0000000000000..7d21af6b33480 --- /dev/null +++ b/aten/src/ATen/cuda/StatelessPhilox4x32.cuh @@ -0,0 +1,61 @@ +// Stateless Philox-4x32 PRNG implementation. +// +// Unlike PhiloxRNGEngine (PhiloxUtils.cuh), this is a pure function: given +// (seed, offset) it returns 4 pseudo-random uint32 values with no mutable +// state. This makes it suitable for use in stateless random APIs. +// +// The Philox-4x32 cipher operates on a 128-bit counter. The full counter +// is (offset_lo, offset_hi, subsequence_lo, subsequence_hi), but we fix +// subsequence=0 so that the entire 128-bit counter space is addressed by +// the 64-bit offset alone. This keeps the API simple and maintains +// cross-device consistency. For example, utilizing thread ID-based subsequence +// numbers and SM-based thread count causes different random values to +// be generated across GPU types. We avoid this situation by always setting +// subsequence=0. + +#pragma once + +#include + +namespace at::cuda { + +__device__ __forceinline__ uint2 mulhilo32(uint32_t a, uint32_t b) { + return {a * b, __umulhi(a, b)}; +} + +__device__ __forceinline__ uint4 philox_round(uint4 ctr, uint2 key) { + constexpr uint32_t kPhiloxSA = 0xD2511F53; + constexpr uint32_t kPhiloxSB = 0xCD9E8D57; + uint2 r0 = mulhilo32(kPhiloxSA, ctr.x); + uint2 r1 = mulhilo32(kPhiloxSB, ctr.z); + return {r1.y ^ ctr.y ^ key.x, r1.x, r0.y ^ ctr.w ^ key.y, r0.x}; +} + +// Stateless Philox-4x32. Returns 4 pseudo-random uint32 values (128 bits) +// determined entirely by (seed, offset). Each unique offset produces a +// distinct 128-bit output. +template +__device__ __forceinline__ uint4 philox_4x32( + uint64_t seed, uint64_t offset) { + uint2 key = { + static_cast(seed), + static_cast(seed >> 32)}; + uint4 ctr = { + static_cast(offset), + static_cast(offset >> 32), + // restrict subsequence=0 + 0, 0}; + + constexpr uint32_t kPhilox10A = 0x9E3779B9; + constexpr uint32_t kPhilox10B = 0xBB67AE85; + + #pragma unroll + for (int i = 0; i < N_ROUNDS - 1; i++) { + ctr = philox_round(ctr, key); + key.x += kPhilox10A; + key.y += kPhilox10B; + } + return philox_round(ctr, key); +} + +} // namespace at::cuda diff --git a/aten/src/ATen/cuda/detail/CUDAHooks.cpp b/aten/src/ATen/cuda/detail/CUDAHooks.cpp index 03a3a97525a43..5f81407b1ac03 100644 --- a/aten/src/ATen/cuda/detail/CUDAHooks.cpp +++ b/aten/src/ATen/cuda/detail/CUDAHooks.cpp @@ -4,7 +4,6 @@ #include #include #include -#include #include #include #include @@ -86,9 +85,6 @@ struct _Initializer { // let's not if we don't need to!) void CUDAHooks::init() const { C10_LOG_API_USAGE_ONCE("aten.init.cuda"); - // Force the update to enable unit testing. This code get executed before unit tests - // have a chance to enable vitals. - at::vitals::VitalsAPI.setVital("CUDA", "used", "true", /* force = */ true); const auto num_devices = c10::cuda::device_count_ensure_non_zero(); c10::cuda::CUDACachingAllocator::init(num_devices); diff --git a/aten/src/ATen/cuda/detail/DeviceThreadHandles.h b/aten/src/ATen/cuda/detail/DeviceThreadHandles.h index 71a344d281d2a..9ec194fc878b6 100644 --- a/aten/src/ATen/cuda/detail/DeviceThreadHandles.h +++ b/aten/src/ATen/cuda/detail/DeviceThreadHandles.h @@ -15,6 +15,7 @@ #pragma once +#include #include #include #include @@ -114,9 +115,45 @@ struct DeviceThreadHandlePool : public std::enable_shared_from_this guard(parent->mutex); + + if(parent->available_handles[device].size() > 0) + { + my_stream_handles[key] = parent->available_handles[device].back(); + parent->available_handles[device].pop_back(); + } + else + { + parent->created_handles[device].emplace_back(true /*create*/); + my_stream_handles[key] = parent->created_handles[device].back().handle; + } + + return my_stream_handles[key]; + } +#endif + private: // Stores the per-device handles currently owned by this thread std::unordered_map my_handles; +#ifdef USE_ROCM + // Stores per-(device, stream) handles for ROCm, where hipblaslt + // requires a unique handle per stream. + std::map, Handle_t> my_stream_handles; +#endif std::weak_ptr weak_parent; @@ -134,6 +171,18 @@ struct DeviceThreadHandlePool : public std::enable_shared_from_thisavailable_handles[d_h.first].push_back(d_h.second); } +#ifdef USE_ROCM + if(!my_stream_handles.empty()) { + auto parent = weak_parent.lock(); + if (!parent) { + return; + } + + std::lock_guard guard(parent->mutex); + for(auto& [key, handle] : my_stream_handles) + parent->available_handles[key.first].push_back(handle); + } +#endif } }; diff --git a/aten/src/ATen/cuda/tunable/GemmHipblaslt.h b/aten/src/ATen/cuda/tunable/GemmHipblaslt.h index 29affa2d21ff1..29c15720f4a66 100644 --- a/aten/src/ATen/cuda/tunable/GemmHipblaslt.h +++ b/aten/src/ATen/cuda/tunable/GemmHipblaslt.h @@ -94,6 +94,16 @@ constexpr hipDataType HipDataTypeFor() { #endif } +template +constexpr hipblasComputeType_t HipBlasComputeTypeFor() { + return HIPBLAS_COMPUTE_32F; +} + +template <> +constexpr hipblasComputeType_t HipBlasComputeTypeFor() { + return HIPBLAS_COMPUTE_64F; +} + template int GetBatchFromParams(const GemmParams* params) { return 1; @@ -175,43 +185,43 @@ int GetStrideCFromParams(const ScaledGemmParams* params) { } template -float GetAlphaFromParams(const GemmParams* params) { +at::opmath_type GetAlphaFromParams(const GemmParams* params) { return params->alpha; } template -float GetAlphaFromParams(const GemmAndBiasParams* params) { +at::opmath_type GetAlphaFromParams(const GemmAndBiasParams* params) { return params->alpha; } template -float GetAlphaFromParams(const GemmStridedBatchedParams* params) { +at::opmath_type GetAlphaFromParams(const GemmStridedBatchedParams* params) { return params->alpha; } template -float GetAlphaFromParams(const ScaledGemmParams* params) { - return 1.0; +at::opmath_type GetAlphaFromParams(const ScaledGemmParams* params) { + return at::opmath_type{1.0}; } template -float GetBetaFromParams(const GemmParams* params) { +at::opmath_type GetBetaFromParams(const GemmParams* params) { return params->beta; } template -float GetBetaFromParams(const GemmAndBiasParams* params) { - return 0.0; +at::opmath_type GetBetaFromParams(const GemmAndBiasParams* params) { + return at::opmath_type{0.0}; } template -float GetBetaFromParams(const GemmStridedBatchedParams* params) { +at::opmath_type GetBetaFromParams(const GemmStridedBatchedParams* params) { return params->beta; } template -float GetBetaFromParams(const ScaledGemmParams* params) { - return 0.0; +at::opmath_type GetBetaFromParams(const ScaledGemmParams* params) { + return at::opmath_type{0.0}; } template @@ -467,8 +477,9 @@ class HipblasltGemmOp : public Callable { TORCH_CHECK(transa_outer == opa && transb_outer == opb, "trans mismatch, shouldn't happen"); - float alpha = GetAlphaFromParams(params); - float beta = GetBetaFromParams(params); + using opmath_t = at::opmath_type; + opmath_t alpha = GetAlphaFromParams(params); + opmath_t beta = GetBetaFromParams(params); hipblasLtMatrixLayout_t mat_a, mat_b, mat_c; if (opa == HIPBLAS_OP_N) { @@ -505,11 +516,14 @@ class HipblasltGemmOp : public Callable { mat_c, HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &stride_c, sizeof(stride_c))); } - hipblasComputeType_t computeType = HIPBLAS_COMPUTE_32F; - if (at::globalContext().float32Precision(at::Float32Backend::CUDA, at::Float32Op::MATMUL) == at::Float32Precision::TF32) { - computeType = HIPBLAS_COMPUTE_32F_FAST_TF32; + hipblasComputeType_t computeType = HipBlasComputeTypeFor(); + if constexpr (std::is_same_v) { + if (at::globalContext().float32Precision(at::Float32Backend::CUDA, at::Float32Op::MATMUL) == at::Float32Precision::TF32) { + computeType = HIPBLAS_COMPUTE_32F_FAST_TF32; + } } - HipBlasLtMatmulDescriptor matmul(computeType, HIP_R_32F); + auto scale_type = HipDataTypeFor(); + HipBlasLtMatmulDescriptor matmul(computeType, scale_type); matmul.setAttribute(HIPBLASLT_MATMUL_DESC_TRANSA, opa); matmul.setAttribute(HIPBLASLT_MATMUL_DESC_TRANSB, opb); @@ -630,9 +644,11 @@ auto GetHipBlasLtTypeStringAndOps() { } #endif - hipblasComputeType_t computeType = HIPBLAS_COMPUTE_32F; - if (at::globalContext().allowTF32CuBLAS()) { - computeType = HIPBLAS_COMPUTE_32F_FAST_TF32; + hipblasComputeType_t computeType = HipBlasComputeTypeFor(); + if constexpr (std::is_same_v) { + if (at::globalContext().float32Precision(at::Float32Backend::CUDA, at::Float32Op::MATMUL) == at::Float32Precision::TF32) { + computeType = HIPBLAS_COMPUTE_32F_FAST_TF32; + } } hipblasLtHandle_t handle; diff --git a/aten/src/ATen/functorch/BatchRulesModules.cpp b/aten/src/ATen/functorch/BatchRulesModules.cpp index 4e0b50c4e3fe7..51cb1d4751d6f 100644 --- a/aten/src/ATen/functorch/BatchRulesModules.cpp +++ b/aten/src/ATen/functorch/BatchRulesModules.cpp @@ -168,13 +168,20 @@ static std::tuple, Tensor, std::optional grid_sample_backward_helper_out( // NOLINTNEXTLINE(performance-unnecessary-value-param) std::tuple bw_out, - int64_t grad_input_out_bdim, - int64_t grad_grid_out_bdim, + int64_t grad_input_bdim, + int64_t grad_grid_bdim, int64_t bdim_size) { auto& [grad_input, grad_grid] = bw_out; - grad_input = reshape_dim_outof(grad_input_out_bdim, bdim_size, grad_input); - grad_grid = reshape_dim_outof(grad_grid_out_bdim, bdim_size, grad_grid); - return std::make_tuple(std::move(grad_input), grad_input_out_bdim, std::move(grad_grid), grad_grid_out_bdim); + std::optional grad_input_bdim_out, grad_grid_bdim_out; + if (grad_input.defined()) { + grad_input = reshape_dim_outof(grad_input_bdim, bdim_size, grad_input); + grad_input_bdim_out = grad_input_bdim; + } + if (grad_grid.defined()) { + grad_grid = reshape_dim_outof(grad_grid_bdim, bdim_size, grad_grid); + grad_grid_bdim_out = grad_grid_bdim; + } + return std::make_tuple(std::move(grad_input), grad_input_bdim_out, std::move(grad_grid), grad_grid_bdim_out); } diff --git a/aten/src/ATen/functorch/BatchRulesNorm.cpp b/aten/src/ATen/functorch/BatchRulesNorm.cpp index 292cf581b9467..affb6ce369f2b 100644 --- a/aten/src/ATen/functorch/BatchRulesNorm.cpp +++ b/aten/src/ATen/functorch/BatchRulesNorm.cpp @@ -351,7 +351,7 @@ static at::Tensor group_norm_backward_no_weight_bias_batch_rule( auto rstd_ = moveBatchDimToFront(rstd, rstd_bdim); const auto bdim_size = get_bdim_size2(grad_out, grad_out_bdim, input, input_bdim); - grad_out_ = ensure_has_bdim(grad_out, grad_out_bdim.has_value(), bdim_size); + grad_out_ = ensure_has_bdim(grad_out_, grad_out_bdim.has_value(), bdim_size); input_ = ensure_has_bdim(input_, input_bdim.has_value(), bdim_size); mean_ = ensure_has_bdim(mean_, mean_bdim.has_value(), bdim_size); rstd_ = ensure_has_bdim(rstd_, rstd_bdim.has_value(), bdim_size); diff --git a/aten/src/ATen/functorch/BatchRulesScatterOps.cpp b/aten/src/ATen/functorch/BatchRulesScatterOps.cpp index 3e75831c071a5..eab17f23ce2bf 100644 --- a/aten/src/ATen/functorch/BatchRulesScatterOps.cpp +++ b/aten/src/ATen/functorch/BatchRulesScatterOps.cpp @@ -774,11 +774,22 @@ std::tuple> scatter_add_batch_rule( self, self_bdim, dim, index, index_bdim, src, src_bdim); } +static void check_scatter_inplace_bdim( + std::optional self_bdim, + std::optional index_bdim, + std::optional src_bdim, + const char* schema_name) { + if (!self_bdim.has_value() && (index_bdim.has_value() || src_bdim.has_value())) { + vmapIncompatibleInplaceError(schema_name); + } +} + std::tuple> scatter_add__batch_rule( const Tensor& self, std::optional self_bdim, int64_t dim, const Tensor& index, std::optional index_bdim, const Tensor& src, std::optional src_bdim) { + check_scatter_inplace_bdim(self_bdim, index_bdim, src_bdim, "scatter_add_"); return scatter_batch_rule(ATEN_FN(scatter_add_), self, self_bdim, dim, index, index_bdim, src, src_bdim); } @@ -811,6 +822,8 @@ std::tuple> scatter_reduce__two_batch_rule( const Tensor& src, std::optional src_bdim, const std::string_view reduce, bool include_self) { + check_scatter_inplace_bdim( + self_bdim, index_bdim, src_bdim, "scatter_reduce_"); return scatter_batch_rule(ATEN_FN2(scatter_reduce_, two), self, self_bdim, dim, index, index_bdim, src, src_bdim, reduce, include_self); } diff --git a/aten/src/ATen/mps/MPSDevice.h b/aten/src/ATen/mps/MPSDevice.h index 3c2dee71aa752..67f65b99c79e0 100644 --- a/aten/src/ATen/mps/MPSDevice.h +++ b/aten/src/ATen/mps/MPSDevice.h @@ -23,6 +23,7 @@ enum class MacOSVersion : uint32_t { MACOS_VER_15_1_PLUS, MACOS_VER_15_2_PLUS, MACOS_VER_26_0_PLUS, + MACOS_VER_26_4_PLUS, }; //----------------------------------------------------------------- diff --git a/aten/src/ATen/mps/MPSDevice.mm b/aten/src/ATen/mps/MPSDevice.mm index 89b54d3d3d047..4e864aeaa73f8 100644 --- a/aten/src/ATen/mps/MPSDevice.mm +++ b/aten/src/ATen/mps/MPSDevice.mm @@ -66,6 +66,7 @@ static inline MTLLanguageVersion getMetalLanguageVersion(const id& de static bool _macos_15_1_plus = is_os_version_at_least(15, 1); static bool _macos_15_2_plus = is_os_version_at_least(15, 2); static bool _macos_26_0_plus = is_os_version_at_least(26, 0); + static bool _macos_26_4_plus = is_os_version_at_least(26, 4); switch (version) { case MacOSVersion::MACOS_VER_14_4_PLUS: @@ -78,6 +79,8 @@ static inline MTLLanguageVersion getMetalLanguageVersion(const id& de return _macos_15_2_plus; case MacOSVersion::MACOS_VER_26_0_PLUS: return _macos_26_0_plus; + case MacOSVersion::MACOS_VER_26_4_PLUS: + return _macos_26_4_plus; default: return false; } diff --git a/aten/src/ATen/mps/MPSGuardImpl.h b/aten/src/ATen/mps/MPSGuardImpl.h index 008a8d57f3df6..368515f233e9b 100644 --- a/aten/src/ATen/mps/MPSGuardImpl.h +++ b/aten/src/ATen/mps/MPSGuardImpl.h @@ -84,6 +84,18 @@ struct TORCH_API MPSGuardImpl final Stream exchangeStream(Stream s) const override { return Stream(Stream::DEFAULT, Device(c10::DeviceType::MPS, 0)); } + DeviceCapability getDeviceCapability(Device /* unused */) const override { + DeviceCapability cap; + cap.capability_data.capability_bits = (1ULL << kIndex_Byte) | + (1ULL << kIndex_Char) | (1ULL << kIndex_Short) | (1ULL << kIndex_Int) | + (1ULL << kIndex_Long) | (1ULL << kIndex_Half) | (1ULL << kIndex_Float) | + (1ULL << kIndex_ComplexHalf) | (1ULL << kIndex_ComplexFloat) | + (1ULL << kIndex_Bool) | (1ULL << kIndex_BFloat16) | + (1ULL << kIndex_UInt32) | (1ULL << kIndex_UInt16) | + (1ULL << kIndex_UInt64); + return cap; + } + DeviceIndex deviceCount() const noexcept override { if (at::hasMPS()) { // TODO: extend it for multi-device case diff --git a/aten/src/ATen/mps/MPSStream.h b/aten/src/ATen/mps/MPSStream.h index b00890b9f5901..f83a387ee846e 100644 --- a/aten/src/ATen/mps/MPSStream.h +++ b/aten/src/ATen/mps/MPSStream.h @@ -73,7 +73,6 @@ class TORCH_API MPSStream { MTLComputeCommandEncoder_t commandEncoder(); void endKernelCoalescing(); void synchronize(SyncType syncType); - void fill(MTLBuffer_t buffer, uint8_t value, size_t length, size_t offset, SyncType syncType = SyncType::NONE); void copy(MTLBuffer_t srcBuffer, MTLBuffer_t dstBuffer, size_t length, diff --git a/aten/src/ATen/mps/MPSStream.mm b/aten/src/ATen/mps/MPSStream.mm index b02d2447ce17f..8cecc94f5926b 100644 --- a/aten/src/ATen/mps/MPSStream.mm +++ b/aten/src/ATen/mps/MPSStream.mm @@ -157,33 +157,6 @@ @interface MPSGraphExecutionDescriptor () }); } -void MPSStream::fill(id buffer, uint8_t value, size_t length, size_t offset, SyncType syncType) { - if (length == 0) { - return; - } - dispatch_sync_with_rethrow(_serialQueue, ^() { - @autoreleasepool { - endKernelCoalescing(); - id blitEncoder = [commandBuffer() blitCommandEncoder]; - - // For some reason fillBufferfor stopped working for length > 4Gb on MacOS 26 - // See https://github.com/pytorch/pytorch/issues/163962 - // Workaround by batching copy commands into 4Gb chunks - constexpr size_t max_copy_size = 0x100000000; // 4GB - size_t bytes_filled = 0; - size_t bytes_remains = length; - while (bytes_remains > 0) { - NSUInteger bytes_to_copy = std::min(max_copy_size, bytes_remains); - [blitEncoder fillBuffer:buffer range:NSMakeRange(offset + bytes_filled, bytes_to_copy) value:value]; - bytes_filled += bytes_to_copy; - bytes_remains -= bytes_to_copy; - } - [blitEncoder endEncoding]; - synchronize(syncType); - } - }); -} - void MPSStream::copy(id srcBuffer, id dstBuffer, size_t length, diff --git a/aten/src/ATen/native/BatchLinearAlgebra.cpp b/aten/src/ATen/native/BatchLinearAlgebra.cpp index cd24d51c62c1c..2de3e24b584b8 100644 --- a/aten/src/ATen/native/BatchLinearAlgebra.cpp +++ b/aten/src/ATen/native/BatchLinearAlgebra.cpp @@ -1843,7 +1843,21 @@ TORCH_IMPL_FUNC(linalg_cholesky_ex_out)(const Tensor& A, cholesky_stub(L.device().type(), L, info, upper); - if (!cpu) { + // On non-CPU devices (MAGMA) the pre-copy doesn't zero the unused triangle, + // so we must clean up after. On macOS, Accelerate's LAPACK writes into the + // unreferenced triangle for matrices larger than its internal block size + // (e.g. n > 64), violating the LAPACK spec which says "not referenced" + // elements are "never read, written to, or otherwise accessed" + // (see https://www.netlib.org/lapack/lug/node121.html). + // We work around this by applying the same cleanup on macOS. + // TODO(https://github.com/pytorch/pytorch/issues/179152): always + // clean up the unused triangle on all platforms. +#if defined(__APPLE__) + constexpr bool needs_triangle_cleanup = true; +#else + const bool needs_triangle_cleanup = !cpu; +#endif + if (needs_triangle_cleanup) { if (upper) { L.triu_(); } else { diff --git a/aten/src/ATen/native/Distance.cpp b/aten/src/ATen/native/Distance.cpp index b7dfef7779ca7..af8d6185d4c26 100644 --- a/aten/src/ATen/native/Distance.cpp +++ b/aten/src/ATen/native/Distance.cpp @@ -103,8 +103,8 @@ static Tensor cdist_impl(const Tensor& x1, const Tensor& x2, const double p, std // See Note [cdist relies on cdist_impl redispatching] // Keep this condition in sync with the condition at the Note if (!(p == 2 && (mode == 1 || (mode == 0 && (r1 > 25 || r2 > 25))))) { - TORCH_CHECK(device1 == kCPU || device1 == kCUDA || device1 == kXPU, "cdist only supports CPU, XPU and CUDA devices, X1 got: ", device1); - TORCH_CHECK(device2 == kCPU || device2 == kCUDA || device2 == kXPU, "cdist only supports CPU, XPU and CUDA devices, X2 got: ", device2); + TORCH_CHECK(device1 == kCPU || device1 == kCUDA || device1 == kXPU || device1 == kPrivateUse1, "cdist only supports CPU, XPU, CUDA and PrivateUse1 devices, X1 got: ", device1); + TORCH_CHECK(device2 == kCPU || device2 == kCUDA || device2 == kXPU || device2 == kPrivateUse1, "cdist only supports CPU, XPU, CUDA and PrivateUse1 devices, X2 got: ", device2); } auto dim1 = x1.dim(); @@ -229,9 +229,9 @@ Tensor _cdist_backward(const Tensor& _grad, const Tensor& _x1, const Tensor& _x2 int64_t n = x1.size(-2); int64_t m = x1.size(-1); auto device1 = x1.device().type(); - TORCH_CHECK(device1 == kCPU || device1 == kCUDA || device1 == kXPU, "_cdist_backward only supports CPU, XPU and CUDA devices, X1 got: ", device1); + TORCH_CHECK(device1 == kCPU || device1 == kCUDA || device1 == kXPU || device1 == kPrivateUse1, "_cdist_backward only supports CPU, XPU, CUDA and PrivateUse1 devices, X1 got: ", device1); auto device2 = x2.device().type(); - TORCH_CHECK(device2 == kCPU || device2 == kCUDA || device2 == kXPU, "_cdist_backward only supports CPU, XPU and CUDA devices, X2 got: ", device2); + TORCH_CHECK(device2 == kCPU || device2 == kCUDA || device2 == kXPU || device2 == kPrivateUse1, "_cdist_backward only supports CPU, XPU, CUDA and PrivateUse1 devices, X2 got: ", device2); Tensor grad_x1 = at::empty({batch_product, n, m}, x1.options(), LEGACY_CONTIGUOUS_MEMORY_FORMAT); @@ -245,7 +245,7 @@ Tensor _cdist_backward(const Tensor& _grad, const Tensor& _x1, const Tensor& _x2 Tensor _pdist_forward(const Tensor& self, const double p) { TORCH_CHECK(self.is_contiguous(), "_pdist_forward requires contiguous input"); auto device = self.device().type(); - TORCH_CHECK(device == kCPU || device == kCUDA || device == kXPU, "_pdist_forward only supports CPU, XPU and CUDA devices, got: ", device); + TORCH_CHECK(device == kCPU || device == kCUDA || device == kXPU || device == kPrivateUse1, "_pdist_forward only supports CPU, XPU, CUDA and PrivateUse1 devices, got: ", device); Tensor result = at::empty({0}, self.options(), LEGACY_CONTIGUOUS_MEMORY_FORMAT); if (self.size(0) <= 1) { result.resize_({0}); @@ -266,7 +266,7 @@ Tensor _pdist_backward(const Tensor& grad, const Tensor& self, const double p, c TORCH_CHECK(self.is_contiguous(), "_pdist_backward requires self to be contiguous"); TORCH_CHECK(pdist.is_contiguous(), "_pdist_backward requires pdist to be contiguous"); auto device = self.device().type(); - TORCH_CHECK(device == kCPU || device == kCUDA || device == kXPU, "_pdist_backward only supports CPU, XPU and CUDA devices, got: ", device); + TORCH_CHECK(device == kCPU || device == kCUDA || device == kXPU || device == kPrivateUse1, "_pdist_backward only supports CPU, XPU, CUDA and PrivateUse1 devices, got: ", device); Tensor result = at::empty_like(self, LEGACY_CONTIGUOUS_MEMORY_FORMAT); pdist_backward_stub(device, result, grad, self, p, pdist); return result; diff --git a/aten/src/ATen/native/Distributions.cpp b/aten/src/ATen/native/Distributions.cpp index 5f34ed9d24c17..6382b1ecaace1 100644 --- a/aten/src/ATen/native/Distributions.cpp +++ b/aten/src/ATen/native/Distributions.cpp @@ -594,7 +594,8 @@ Tensor& multinomial_out(const Tensor& self, // https://github.com/pytorch/pytorch/issues/11931#issuecomment-625882503 if (!with_replacement || n_sample == 1) { // Sanity checks on `self`. - auto is_valid = ((self.max() < INFINITY) & (self.min() >= 0)); + auto [self_min, self_max] = self.aminmax(); + auto is_valid = ((self_max < INFINITY) & (self_min >= 0)); at::_assert_async(is_valid, "probability tensor contains either `inf`, `nan` or element < 0"); at::Tensor zero_prob_condition; if (self.dim() == 1){ diff --git a/aten/src/ATen/native/ForeachOpsKernels.cpp b/aten/src/ATen/native/ForeachOpsKernels.cpp index d411a68b037e0..2fc3c81ce0eb1 100644 --- a/aten/src/ATen/native/ForeachOpsKernels.cpp +++ b/aten/src/ATen/native/ForeachOpsKernels.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -364,6 +365,21 @@ FOREACH_BINARY_OP_LIST(div) FOREACH_BINARY_OP_LIST(clamp_min) FOREACH_BINARY_OP_LIST(clamp_max) FOREACH_BINARY_OP_LIST(pow) + +// _foreach_clone +std::vector foreach_tensor_clone_slow( + TensorList self, + std::optional memory_format) { + check_foreach_api_restrictions(self); + + std::vector ret{}; + ret.reserve(self.size()); + for (const auto& t : self) { + ret.emplace_back(t.clone(memory_format)); + } + return ret; +} + // _foreach_copy_ void foreach_tensor_copy_list_kernel_slow_( TensorList self, diff --git a/aten/src/ATen/native/ForeachUtils.h b/aten/src/ATen/native/ForeachUtils.h index f0dce20a6eff4..d68d770f215c0 100644 --- a/aten/src/ATen/native/ForeachUtils.h +++ b/aten/src/ATen/native/ForeachUtils.h @@ -93,22 +93,27 @@ inline void check_foreach_api_restrictions( // same device and dtype. inline bool _check_tensors_share_device_and_dtype( ArrayRef tensorLists, - const bool skip_dtype_check = false) { + const bool skip_cross_list_dtype_check = false) { const auto expected_dtype = tensorLists[0][0].dtype(); const auto expected_device = tensorLists[0][0].device(); - auto is_tensor_okay = [&](const Tensor& tensor) { - return (skip_dtype_check || tensor.dtype() == expected_dtype) && - tensor.device() == expected_device && tensor.layout() == at::kStrided && - tensor.is_non_overlapping_and_dense(); - }; - return std::all_of( tensorLists.cbegin(), tensorLists.cend(), [&](const TensorList& tensorList) { + if (tensorList.empty()) { + return true; + } + const auto list_dtype = tensorList[0].dtype(); return std::all_of( - tensorList.cbegin(), tensorList.cend(), is_tensor_okay); + tensorList.cbegin(), tensorList.cend(), [&](const Tensor& tensor) { + return tensor.device() == expected_device && + tensor.layout() == at::kStrided && + tensor.is_non_overlapping_and_dense() && + tensor.dtype() == list_dtype && + (skip_cross_list_dtype_check || + tensor.dtype() == expected_dtype); + }); }); } @@ -192,8 +197,10 @@ inline bool _check_tensors_do_type_promotion_with_scalars( inline bool check_fast_path_restrictions( ArrayRef tensorLists, ArrayRef scalarList = {}, - bool does_op_promote_integer_inputs_to_float = false) { - return _check_tensors_share_device_and_dtype(tensorLists) && + bool does_op_promote_integer_inputs_to_float = false, + bool skip_cross_list_dtype_check = false) { + return _check_tensors_share_device_and_dtype( + tensorLists, skip_cross_list_dtype_check) && _check_tensors_share_sizes_and_strides(tensorLists) && _check_tensors_do_type_promotion_with_scalars( tensorLists[0], @@ -318,9 +325,17 @@ inline FlatMap _group_tensors_by_first_tensors_device_and_dtype( const auto s = tensor->scalar_type(); const auto d = tensor->device(); // Note: `step` or `state_step` is float32 by default. + // BFloat16 is allowed here for mixed-precision optimizer + // states (e.g. fp32 params with bf16 exp_avg/exp_avg_sq). + // Currently only BFloat16 is permitted because it is the + // only low-precision state dtype validated end-to-end in + // large-scale training (e.g. DeepSeek-V3 671B). + // TBD: make the set of allowed extra dtypes configurable + // per optimizer so this function stays dtype-agnostic. if (key.first == d) { return key.second == s || s == at::ScalarType::Float || - s == at::ScalarType::Double; + s == at::ScalarType::Double || + s == at::ScalarType::BFloat16; } else if (d.is_cpu()) { // note(crcrpar): There are some test cases (e.g. // TestOptim::test_adam) where state_steps are on CPU and the @@ -333,7 +348,9 @@ inline FlatMap _group_tensors_by_first_tensors_device_and_dtype( } } }), - "Tensors of the same index must be on the same device and the same dtype except `step` tensors that can be CPU and float32/64 notwithstanding"); + "Tensors of the same index must be on the same device and the same dtype " + "except `step` tensors that can be CPU and float32/64, and optimizer " + "states that can be bfloat16 for mixed-precision training"); grouped_tensors_with_indices.try_emplace( key, TensorsAndIndicesT{ diff --git a/aten/src/ATen/native/Math.h b/aten/src/ATen/native/Math.h index ed45ff6ebe803..7d9e73779a4bd 100644 --- a/aten/src/ATen/native/Math.h +++ b/aten/src/ATen/native/Math.h @@ -2281,7 +2281,7 @@ inline C10_HOST_DEVICE T airy_ai_forward(T x) { int domain_flag = 0; - T ai; + T ai = T(0.0); if (std::isinf(x)) { return std::numeric_limits::quiet_NaN(); diff --git a/aten/src/ATen/native/Normalization.cpp b/aten/src/ATen/native/Normalization.cpp index 22e2d4b75e950..71dab6aed8955 100644 --- a/aten/src/ATen/native/Normalization.cpp +++ b/aten/src/ATen/native/Normalization.cpp @@ -118,14 +118,14 @@ struct Var { } }; -static inline bool is_contiguous(const Tensor& t) { +static bool is_contiguous_in_any_format(const Tensor& t) { return t.is_contiguous() || t.is_contiguous(at::MemoryFormat::ChannelsLast) || t.is_contiguous(at::MemoryFormat::ChannelsLast3d); } // For some ambiguous cases, it is possible a channels last contiguous Tensor has // `suggest_memory_format` of Contiguous. // See https://github.com/pytorch/pytorch/issues/63224 for details. -static inline MemoryFormat suggest_memory_format_contig(const Tensor& t) { +static MemoryFormat suggest_memory_format_contig(const Tensor& t) { return t.is_contiguous() ? at::MemoryFormat::Contiguous : (t.is_contiguous(at::MemoryFormat::ChannelsLast3d) ? at::MemoryFormat::ChannelsLast3d : at::MemoryFormat::ChannelsLast); @@ -138,8 +138,8 @@ static std::tuple batch_norm_cpu_transform_input_template( const Tensor& running_mean /* optional */, const Tensor& running_var /* optional */, bool train, double eps, Tensor& output) { - bool all_contiguous = is_contiguous(input) - && is_contiguous(output) + bool all_contiguous = is_contiguous_in_any_format(input) + && is_contiguous_in_any_format(output) && (!weight.defined() || weight.is_contiguous()) && (!bias.defined() || bias.is_contiguous()) && running_mean.is_contiguous() @@ -207,7 +207,7 @@ static std::tuple batch_norm_cpu_update_stats_template( TORCH_CHECK(input.numel() != 0, "input tensor must have at least one element, but got input_sizes = ", input.sizes()); int64_t n = input.numel() / n_input; - bool all_contiguous = is_contiguous(input); + bool all_contiguous = is_contiguous_in_any_format(input); constexpr bool mixed_type = !std::is_same_v; // Using float data type for Half _var_sum in batchnorm stats updating on CPU // to avoid _var_sum overflow since the representation range of Half is small. @@ -300,7 +300,7 @@ static std::tuple batch_norm_cpu_update_stats_template( constexpr bool mixed_type = !std::is_same_v; const auto dtype = mixed_type ? kFloat : input.scalar_type(); - Tensor save_mean = is_contiguous(input) ? at::empty({n_input}, input.options().dtype(dtype)) : at::mean(input, /*dim=*/reduce_dims, /*keepdim=*/false, dtype); + Tensor save_mean = is_contiguous_in_any_format(input) ? at::empty({n_input}, input.options().dtype(dtype)) : at::mean(input, /*dim=*/reduce_dims, /*keepdim=*/false, dtype); Tensor save_var_transform = at::empty({n_input}, input.options().dtype(dtype)); return batch_norm_cpu_update_stats_template(input, running_mean, running_var, momentum, eps, save_mean, save_var_transform); } @@ -331,8 +331,8 @@ static std::tuple batch_norm_backward_cpu_template( // since we are directly manipulating pointers in contiguous path, // need to make sure input and grad_out have the same memory format. - bool all_contiguous = is_contiguous(input) - && is_contiguous(grad_out_) + bool all_contiguous = is_contiguous_in_any_format(input) + && is_contiguous_in_any_format(grad_out_) && input.suggest_memory_format() == grad_out_.suggest_memory_format(); if (all_contiguous) { @@ -863,7 +863,7 @@ std::tuple batch_norm_cpu(const Tensor& self, const std: checkBackend("batch_norm_cpu", {self, weight, bias, running_mean, running_var}, Backend::CPU); // Prepare output tensor - const bool all_contiguous = is_contiguous(self) + const bool all_contiguous = is_contiguous_in_any_format(self) && (!weight.defined() || weight.is_contiguous()) && (!bias.defined() || bias.is_contiguous()) && running_mean.is_contiguous() @@ -885,7 +885,7 @@ std::tuple batch_norm_cpu(const Tensor& self, const std: save_mean = at::empty({0}, self.options().dtype(kFloat)); save_var = at::empty({0}, self.options().dtype(kFloat)); } else { - save_mean = is_contiguous(self) ? at::empty({self.size(1)}, self.options().dtype(kFloat)) : at::mean(self, /*dim=*/reduce_dims, /*keepdim=*/false, kFloat); + save_mean = is_contiguous_in_any_format(self) ? at::empty({self.size(1)}, self.options().dtype(kFloat)) : at::mean(self, /*dim=*/reduce_dims, /*keepdim=*/false, kFloat); save_var = at::empty({self.size(1)}, self.options().dtype(kFloat)); } } else { @@ -893,7 +893,7 @@ std::tuple batch_norm_cpu(const Tensor& self, const std: save_mean = at::empty({0}, self.options()); save_var = at::empty({0}, self.options()); } else { - save_mean = is_contiguous(self) ? at::empty({self.size(1)}, self.options()) : at::mean(self, /*dim=*/reduce_dims, /*keepdim=*/false); + save_mean = is_contiguous_in_any_format(self) ? at::empty({self.size(1)}, self.options()) : at::mean(self, /*dim=*/reduce_dims, /*keepdim=*/false); save_var = at::empty({self.size(1)}, self.options()); } } diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp index 68dadf218945f..d5dc9ba1eeb46 100644 --- a/aten/src/ATen/native/ReduceOps.cpp +++ b/aten/src/ATen/native/ReduceOps.cpp @@ -117,6 +117,7 @@ #include #include #include +#include #include #include #endif @@ -645,7 +646,12 @@ Tensor cumprod_backward(const Tensor& grad, const Tensor& input, int64_t dim, co // O(n) implementation. The derivative of this implementation is _not_ // the second derivative of cumprod. As such, we fallback to a less efficient // O(n^2) implementation when at::GradMode::is_enabled(). - if (!at::GradMode::is_enabled() && !are_inputs_tensors_sublcass) { + // + // NOTE: We use at::where instead of masked_scatter_/masked_select to make + // this path composite compliant for tensor subclasses (e.g., FakeTensors + // used by torch.compile). masked_select has dynamic output shape which + // causes issues with tracing. See https://github.com/pytorch/pytorch/issues/136263 + if (!at::GradMode::is_enabled()) { // n.b. This could probably be implemented much faster with a kernel // From here on we need to use some mask gymnastics to @@ -669,9 +675,16 @@ Tensor cumprod_backward(const Tensor& grad, const Tensor& input, int64_t dim, co // case k < z1 // select everything before the first zero [0, z1) auto mask = cumsum == 0; - // equiv to grad_input[mask] = deriv[grad] - grad_input.masked_scatter_(mask, - reversed_cumsum(w.masked_fill(~mask, 0.), dim).div_(input_conj).masked_select(mask)); + // Compute gradient for positions before the first zero + // Using at::where instead of masked_scatter_ for composite compliance + auto grad_before_first_zero = reversed_cumsum(w.masked_fill(~mask, 0.), dim); + if (!are_inputs_tensors_sublcass) { + grad_before_first_zero = grad_before_first_zero.div_(input_conj); + } else { + grad_before_first_zero = grad_before_first_zero.div(input_conj); + } + grad_input = at::where(mask, grad_before_first_zero, grad_input); + // select everything from the first zero to the second zero [z1, z2) mask = cumsum == 1; @@ -693,13 +706,21 @@ Tensor cumprod_backward(const Tensor& grad, const Tensor& input, int64_t dim, co // dy_j / dx_z1 = sum(cumprod(input[z1+1:z2] * grad[z1+1:z2])) * prod(output[z1-1]) // relu_() necessary as gather does not support negative indices // finally, we do grad_input[z1] = dy_j / dx_z1 - grad_input.masked_scatter_(first_zero_mask, - input_conj.masked_fill(~mask, 1.).cumprod(dim) - .mul_(grad.masked_fill(cumsum != 1, 0.)) - .sum(dim, /*keepdim*/true) - .mul_(at::gather(output_conj, dim, (first_zero_index - 1).relu_()) - .masked_fill_(first_zero_index == 0, 1.)) - .masked_select(first_zero_mask)); + // Using at::where instead of masked_scatter_ for composite compliance + auto grad_at_first_zero = input_conj.masked_fill(~mask, 1.).cumprod(dim); + const auto grad_masked = grad.masked_fill(cumsum != 1, 0.); + const auto output_before_zero = at::gather(output_conj, dim, (first_zero_index - 1).relu_()) + .masked_fill_(first_zero_index == 0, 1.); + if (!are_inputs_tensors_sublcass) { + grad_at_first_zero = grad_at_first_zero.mul_(grad_masked) + .sum(dim, /*keepdim*/true) + .mul_(output_before_zero); + } else { + grad_at_first_zero = grad_at_first_zero.mul(grad_masked) + .sum(dim, /*keepdim*/true) + .mul(output_before_zero); + } + grad_input = at::where(first_zero_mask, grad_at_first_zero, grad_input); return grad_input; } else { // GradMode::enabled() /* @@ -1864,8 +1885,8 @@ static Tensor& std_var_out( const char* fname, Tensor& result, const Tensor& self, at::OptionalIntArrayRef dim, const std::optional& correction_opt, bool keepdim, bool take_sqrt) { - TORCH_CHECK(self.device().is_cpu() || self.device().is_cuda() || self.device().is_xpu(), - "std and var supports tensors on a CPU, CUDA, or XPU device only, but got: ", + TORCH_CHECK(self.device().is_cpu() || self.device().is_cuda() || self.device().is_xpu() || self.device().is_privateuseone(), + "std and var supports tensors on a CPU, CUDA, XPU or PrivateUse1 device only, but got: ", self.device().type()); TORCH_CHECK(self.layout() == Layout::Strided, "std and var only supports strided layout, got: ", self.layout()); @@ -1937,8 +1958,8 @@ static std::tuple std_var_mean_out( at::OptionalIntArrayRef dim, const std::optional& correction_opt, bool keepdim, bool take_sqrt) { AT_ASSERT(result1.defined() && result2.defined()); - TORCH_CHECK(self.device().is_cpu() || self.is_cuda() || self.is_xpu(), - fname, " supports tensors on a CPU, CUDA, or XPU device only, got: ", + TORCH_CHECK(self.device().is_cpu() || self.is_cuda() || self.is_xpu() || self.is_privateuseone(), + fname, " supports tensors on a CPU, CUDA, XPU or PrivateUse1 device only, got: ", self.device().type()); TORCH_CHECK(self.layout() == Layout::Strided, fname, " only supports strided layout, got: ", self.layout()); @@ -2351,10 +2372,14 @@ bool cpu_equal(const Tensor& self, const Tensor& other) { Tensor value_selecting_reduction_backward_symint(const Tensor& grad, int64_t dim, const Tensor& indices, c10::SymIntArrayRef sizes, bool keepdim) { auto inplace_scatter_if_not_tensor_subclass = [&](const Tensor& grad_out, const Tensor& indices_) { - auto grad_in = at::zeros_symint(sizes, grad_out.options()); if (areAnyTensorSubclassLike({grad, indices})) { + // Use new_zeros_symint so that tensor subclasses (e.g. DTensor) + // can intercept the zeros creation through dispatch, ensuring + // the result has matching subclass type for subsequent scatter. + auto grad_in = grad_out.new_zeros_symint(sizes); return grad_in.scatter(dim, indices_, grad_out); } + auto grad_in = at::zeros_symint(sizes, grad_out.options()); return grad_in.scatter_(dim, indices_, grad_out); }; diff --git a/aten/src/ATen/native/ScaledBlas.cpp b/aten/src/ATen/native/ScaledBlas.cpp index d58593bfe018f..0b384d7226b73 100644 --- a/aten/src/ATen/native/ScaledBlas.cpp +++ b/aten/src/ATen/native/ScaledBlas.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #if !defined(__s390x__) && !defined(__powerpc__) #include #endif @@ -32,12 +33,102 @@ #include #include #include +#include #include #include #endif namespace at::native { +using at::blas::ScalingType; +using at::blas::SwizzleType; + +namespace { + +void invalid_scaling_config( + const at::Tensor& mat_a, + const at::Tensor& mat_b, + const std::optional& scale_a, + const std::optional& scale_b) { + std::stringstream exception_ss; + exception_ss << "Invalid scaling configuration.\n"; + exception_ss << "- For TensorWise scaling, a and b should be float8, scales should be float and singletons.\n"; + exception_ss << "- For RowWise scaling, a and b should be float8, scales should be float, scale_a should be ("; + exception_ss << mat_a.size(0) << ", 1) and scale_b should be (1, " << mat_b.size(1) << "), and both should be contiguous.\n"; + exception_ss << "Got mat_a.dtype()=" << mat_a.scalar_type(); + if (scale_a.has_value()) { + exception_ss << ", scale_a.dtype()=" << scale_a.value().scalar_type() << ", scale_a.size()=" << scale_a.value().sizes() << ", scale_a.stride()=" << scale_a.value().strides(); + } + else { + exception_ss << ", scale_a=None"; + } + exception_ss << ", mat_b.dtype()=" << mat_b.scalar_type(); + if (scale_b.has_value()) { + exception_ss << ", scale_b.dtype()=" << scale_b.value().scalar_type() << ", scale_b.size()=" << scale_b.value().sizes() << " and scale_b.stride()=" << scale_b.value().strides(); + } + else { + exception_ss << " and scale_b=None"; + } + + TORCH_CHECK_VALUE(false, exception_ss.str()); +} + +/* + * Scaling Type Determination: + * --------------------------- + * Conditions and corresponding Scaling Types: + * + * - If scale.numel() == 1: + * - Returns TensorWise. + * + * - Else if scale.dim() == 2 && scale.size(0) == outer_dim && scale.size(1) == + * 1: + * - Returns RowWise. + * + * - Otherwise: + * - Returns Error. + */ + +bool is_tensorwise_scaling(const at::Tensor& t, const at::Tensor& scale) { + return at::isFloat8Type(t.scalar_type()) && + scale.scalar_type() == at::kFloat && scale.numel() == 1; +} + +bool is_rowwise_scaling(const at::Tensor& t, const at::Tensor& scale) { + return ( + at::isFloat8Type(t.scalar_type()) && scale.scalar_type() == at::kFloat && + scale.dim() == 2 && scale.size(0) == t.size(0) && scale.size(1) == 1 && + scale.is_contiguous()); +} + +bool is_desired_scaling( + const at::Tensor& t, + const at::Tensor& scale, + ScalingType desired_scaling) { + auto result = desired_scaling == ScalingType::TensorWise + ? is_tensorwise_scaling(t, scale) + : is_rowwise_scaling(t, scale); + return result; +} + +std::pair get_joint_scaling( + std::initializer_list> options, + const at::Tensor& mat_a, + const at::Tensor& mat_b, + const at::Tensor& scale_a, + const at::Tensor& scale_b) { + for (auto [lhs, rhs] : options) { + if (is_desired_scaling(mat_a, scale_a, lhs) && + is_desired_scaling(mat_b.t(), scale_b.t(), rhs)) { + return {lhs, rhs}; + } + } + invalid_scaling_config(mat_a, mat_b, scale_a, scale_b); + return {}; +} + +} // namespace + static Tensor& _scaled_mm_out_cpu_emulated(const Tensor& mat1, const Tensor& mat2, const Tensor& scale_a, @@ -46,14 +137,70 @@ _scaled_mm_out_cpu_emulated(const Tensor& mat1, const Tensor& mat2, const std::optional& scale_result, std::optional out_dtype, bool use_fast_accum, + const scaled::ScaledGemmImplementation gemm_impl, Tensor& out) { TORCH_CHECK(mat1.dim() == 2, "mat1 must be a matrix"); TORCH_CHECK(mat2.dim() == 2, "mat2 must be a matrix"); TORCH_CHECK( - mat1.sizes()[1] == mat2.sizes()[0], "mat1 and mat2 shapes cannot be multiplied (", + mat1.sizes()[1] == mat2.sizes()[0], "mat_a and mat_b shapes cannot be multiplied (", mat1.sizes()[0], "x", mat1.sizes()[1], " and ", mat2.sizes()[0], "x", mat2.sizes()[1], ")"); - TORCH_INTERNAL_ASSERT((scale_a.numel() == 1 && scale_b.numel() == 1), "Now _scaled_mm only supports per-tensor scaling for CPU backend."); + TORCH_CHECK_VALUE(gemm_impl == scaled::ScaledGemmImplementation::TENSORWISE_TENSORWISE || + gemm_impl == scaled::ScaledGemmImplementation::ROWWISE_ROWWISE, "Unsupported scaling implementation"); + + if (gemm_impl == scaled::ScaledGemmImplementation::TENSORWISE_TENSORWISE) { + // Restrictions: + // A, B are FP8, scales are fp32 + TORCH_CHECK_VALUE( + isFloat8Type(mat1.scalar_type()) && isFloat8Type(mat2.scalar_type()), + "mat1 and mat2 must be fp8 types, got: ", + mat1.scalar_type(), + mat2.scalar_type()); + TORCH_CHECK_VALUE( + scale_a.numel() == 1 && scale_a.scalar_type() == kFloat, + "scale_a must have 1 Float element"); + TORCH_CHECK_VALUE( + scale_b.numel() == 1 && scale_b.scalar_type() == kFloat, + "scale_b must have 1 Float element"); + } + else { + // Restrictions: + // A, B are FP8, scales are fp32, shape M/N for A/B + TORCH_CHECK_VALUE( + isFloat8Type(mat1.scalar_type()) && isFloat8Type(mat2.scalar_type()), + "mat1 and mat_b must be fp8 types, got: ", + mat1.scalar_type(), + mat2.scalar_type()); + TORCH_CHECK_VALUE( + scale_a.size(0) == mat1.size(0) && scale_a.size(1) == 1, + "scale_a must have shape [", + mat1.size(0), + ", 1], got [", + scale_a.sizes(), + "]"); + TORCH_CHECK_VALUE( + scale_a.numel() == mat1.size(0) && scale_a.scalar_type() == kFloat, + "scale_a must have ", + mat1.size(0), + " Float elements, got ", + scale_a.numel()); + TORCH_CHECK_VALUE( + scale_b.numel() == mat2.size(1) && scale_b.scalar_type() == kFloat, + "scale_b must have ", + mat2.size(1), + " Float elements, got ", + scale_b.numel()); + + TORCH_CHECK_VALUE( + scale_a.stride(1) == 1, + "expected scale_a.stride(1) to be 1, but got ", + scale_a.stride(1)); + TORCH_CHECK_VALUE( + scale_b.stride(1) == 1, + "expected scale_b.stride(1) to be 1, but got ", + scale_b.stride(1)); + } + TORCH_CHECK( !scale_result || (scale_result->numel() == 1 && scale_result->scalar_type() == kFloat), @@ -62,32 +209,41 @@ _scaled_mm_out_cpu_emulated(const Tensor& mat1, const Tensor& mat2, " but got ", bias->numel()); // Check types - TORCH_CHECK(!out_dtype || *out_dtype == out.scalar_type(), "out_dtype must match output matrix type"); TORCH_CHECK(isFloat8Type(mat1.scalar_type()), "Expected mat1 to be Float8 matrix got ", mat1.scalar_type()); TORCH_CHECK(isFloat8Type(mat2.scalar_type()), "Expected mat2 to be Float8 matrix got ", mat2.scalar_type()); - auto mat1_c = mat1.contiguous(); - auto mat2_c = mat2.contiguous(); - IntArrayRef mat1_sizes = mat1_c.sizes(); - IntArrayRef mat2_sizes = mat2_c.sizes(); + auto mat1_cont = mat1.contiguous(); + auto mat2_cont = mat2.contiguous(); + IntArrayRef mat1_sizes = mat1_cont.sizes(); + IntArrayRef mat2_sizes = mat2_cont.sizes(); at::native::resize_output(out, {mat1_sizes[0], mat2_sizes[1]}); - float input_scale = scale_a.item(); - float weight_scale = scale_b.item(); + const auto out_dtype_ = out_dtype.value_or(c10::ScalarType::BFloat16); + float output_scale = 1.0f; if (scale_result.has_value() && - (*out_dtype == ScalarType::Float8_e4m3fn || - *out_dtype == ScalarType::Float8_e5m2)) { + (out_dtype_ == ScalarType::Float8_e4m3fn || + out_dtype_ == ScalarType::Float8_e5m2)) { output_scale = scale_result.value().item(); } - auto fp32_mat1 = at::mul(mat1.to(kFloat), input_scale); - auto fp32_mat2 = at::mul(mat2_c.to(kFloat), weight_scale); - auto out_tmp = at::matmul(fp32_mat1, fp32_mat2); + + at::Tensor fp32_mat_a; + at::Tensor fp32_mat_b; + if (gemm_impl == scaled::ScaledGemmImplementation::TENSORWISE_TENSORWISE) { + fp32_mat_a = at::mul(mat1_cont.to(kFloat), scale_a.item()); + fp32_mat_b = at::mul(mat2_cont.to(kFloat), scale_b.item()); + } + else { + fp32_mat_a = at::mul(mat1_cont.to(kFloat), scale_a); + fp32_mat_b = at::mul(mat2_cont.to(kFloat), scale_b); + } + + auto out_tmp = at::matmul(fp32_mat_a, fp32_mat_b); if (bias) { out_tmp.add_(bias.value()); } - if (*out_dtype == ScalarType::Float8_e4m3fn || - *out_dtype == ScalarType::Float8_e5m2) { + if (out_dtype_ == ScalarType::Float8_e4m3fn || + out_dtype_ == ScalarType::Float8_e5m2) { out_tmp = at::mul(out_tmp, 1 / output_scale); } out_tmp = out_tmp.to(out.scalar_type()); @@ -105,7 +261,7 @@ _scaled_mm_out_cpu(const Tensor& mat1, const Tensor& mat2, bool use_fast_accum, Tensor& out) { #if AT_MKLDNN_ENABLED() && !defined(__powerpc__) - if (at::globalContext().userEnabledMkldnn()) { + if (at::globalContext().userEnabledMkldnn() && scale_a.numel() == 1 && scale_b.numel() == 1) { bool mixed_dtype = mat1.scalar_type() != mat2.scalar_type(); if ((!mixed_dtype && cpuinfo_has_x86_amx_int8()) || (mixed_dtype && cpuinfo_has_x86_amx_fp16())) { @@ -123,7 +279,30 @@ _scaled_mm_out_cpu(const Tensor& mat1, const Tensor& mat2, } #endif { - return _scaled_mm_out_cpu_emulated(mat1, mat2, scale_a, scale_b, bias, scale_result, out_dtype, use_fast_accum, out); + TORCH_CHECK_VALUE( + !out_dtype || *out_dtype == out.scalar_type(), + "out_dtype must match output matrix type"); + + const auto [scaling_choice_a, scaling_choice_b] = get_joint_scaling( + { + std::make_pair(ScalingType::TensorWise, ScalingType::TensorWise), + std::make_pair(ScalingType::RowWise, ScalingType::RowWise), + }, + mat1, + mat2, + scale_a, + scale_b); + + scaled::ScaledGemmImplementation gemm_impl{scaled::ScaledGemmImplementation::NONE}; + + if (scaling_choice_a == ScalingType::TensorWise && scaling_choice_b == ScalingType::TensorWise) { + gemm_impl = scaled::ScaledGemmImplementation::TENSORWISE_TENSORWISE; + } + else if (scaling_choice_a == ScalingType::RowWise && scaling_choice_b == ScalingType::RowWise) { + gemm_impl = scaled::ScaledGemmImplementation::ROWWISE_ROWWISE; + } + + return _scaled_mm_out_cpu_emulated(mat1, mat2, scale_a, scale_b, bias, scale_result, out_dtype, use_fast_accum, gemm_impl, out); } } @@ -140,6 +319,201 @@ _scaled_mm_cpu(const Tensor& mat_a, const Tensor& mat_b, return _scaled_mm_out_cpu(mat_a, mat_b, scale_a, scale_b, bias, scale_result, out_dtype, use_fast_accum, out); } +using acceptance_fn = std::function&, + ArrayRef&, + c10::ScalarType, + std::vector&, + ArrayRef&)>; + +namespace scaled_blas = at::native::scaled; +using scaled_blas::convert_int_to_enum; +using scaled_blas::ScaledGemmImplementation; + +std::array, 2> + scale_kernel_dispatch = {{ + {"tensorwise_tensorwise", + scaled_blas::check_tensorwise_recipe, + ScaledGemmImplementation::TENSORWISE_TENSORWISE}, + {"rowwise_rowwise", + scaled_blas::check_rowwise_recipe, + ScaledGemmImplementation::ROWWISE_ROWWISE}, + + }}; + +Tensor& _scaled_mm_cpu_v2_out( + const Tensor& mat_a, + const Tensor& mat_b, + ArrayRef scale_a, + IntArrayRef scale_recipe_a, + IntArrayRef swizzle_a, + ArrayRef scale_b, + IntArrayRef scale_recipe_b, + IntArrayRef swizzle_b, + const std::optional& bias, + const std::optional out_dtype, + IntArrayRef contraction_dim, + bool use_fast_accum, + Tensor& out) { + TORCH_CHECK_VALUE(mat_a.dim() == 2, "mat_a must be a matrix"); + TORCH_CHECK_VALUE(mat_b.dim() == 2, "mat_b must be a matrix"); + + // If any of M, K, N is 0 - return early (the tensorwise/rowwise float8 gemm kernels + // do not support this case). + if (mat_a.size(0) == 0 || mat_a.size(1) == 0 || mat_b.size(1) == 0) { + // `out` was created with `at::empty`. In the case where we are multiplying + // MxK by KxN and K is the zero dim, we need to initialize here to properly + // return a tensor of zeros. + at::native::resize_output(out, {mat_a.size(0), mat_b.size(1)}); + if (mat_a.size(1) == 0) { + out.zero_(); + } + + return out; + } + + // Check if the input matrix sizes can be multiplied + // - if optional contraction dims are provided, use those + // -- mostly for < 1B formats (i.e. nvfp4x2) where cheap .t() is not available. + if (contraction_dim.size() > 0) { + TORCH_CHECK_VALUE(contraction_dim.size() == 2, "contraction_dim must have exactly 2 elements"); + auto mat_a_dim = contraction_dim[0]; + auto mat_b_dim = contraction_dim[1]; + TORCH_CHECK_VALUE( + mat_a.size(mat_a_dim) == mat_b.size(mat_b_dim), "mat_a and mat_b shapes cannot be multiplied (", + mat_a.size(0), "x", mat_a.size(1), " and ", mat_b.size(0), "x", mat_b.size(1), ") ", + "with contraction dims mat_a: ", mat_a_dim, ", mat_b: ", mat_b_dim); + } else { + TORCH_CHECK_VALUE( + mat_a.size(1) == mat_b.size(0), "mat_a and mat_b shapes cannot be multiplied (", + mat_a.size(0), "x", mat_a.size(1), " and ", mat_b.size(0), "x", mat_b.size(1), ")"); + } + + TORCH_CHECK_VALUE( + !bias || bias->numel() == mat_b.sizes()[1], + "Bias must be size ", + mat_b.sizes()[1], + " but got ", + bias->numel()); + + TORCH_CHECK_VALUE( + !out_dtype || *out_dtype == out.scalar_type(), + "out_dtype must match output matrix type"); + + if (bias) { + TORCH_CHECK_VALUE( + bias->scalar_type() == kFloat || + bias->scalar_type() == c10::ScalarType::BFloat16 || + bias->scalar_type() == c10::ScalarType::Half, + "Bias must be Float32 or BFloat16 or Half, but got ", + bias->scalar_type()); + } + + // Align with CUDA's default out to be bf16 + const auto out_dtype_ = out_dtype.value_or(c10::ScalarType::BFloat16); + + // Conversion of implicitly-defined enums to explicit + auto scale_recipe_a_enum = convert_int_to_enum(scale_recipe_a); + auto swizzle_a_enum = convert_int_to_enum(swizzle_a); + auto scale_recipe_b_enum = convert_int_to_enum(scale_recipe_b); + auto swizzle_b_enum = convert_int_to_enum(swizzle_b); + + if (!swizzle_a_enum.empty() && !swizzle_b_enum.empty()) { + TORCH_CHECK_VALUE( + swizzle_a_enum[0] == at::blas::SwizzleType::NO_SWIZZLE && + swizzle_b_enum[0] == at::blas::SwizzleType::NO_SWIZZLE, + "CPU does not support swizzle."); + } + + // at this point we can start working out what we want to be doing + // Try to do as few steps as possible. + // NOTE: support is deliberately sparse, can explicitly enumerate all + // combinations allowed. Do this via a list of defined (name, acceptance, + // concrete_impl) tuples. + bool found_impl = false; + ScaledGemmImplementation gemm_impl = ScaledGemmImplementation::NONE; + + for (const auto& fn_entry : scale_kernel_dispatch) { + const auto [name, accept_fn, scaled_gemm_impl] = fn_entry; + const bool ok = accept_fn( + mat_a.scalar_type(), + scale_recipe_a_enum, + scale_a, + mat_b.scalar_type(), + scale_recipe_b_enum, + scale_b); + + if (ok) { + gemm_impl = scaled_gemm_impl; + found_impl = true; + break; + } + } + + if (!found_impl) { + const std::optional scale_a_opt = scale_a.empty() ? std::optional{std::nullopt} : std::optional{scale_a[0]}; + const std::optional scale_b_opt = scale_b.empty() ? std::optional{std::nullopt} : std::optional{scale_b[0]}; + + invalid_scaling_config(mat_a, mat_b, scale_a_opt, scale_b_opt); + } + + at::native::resize_output(out, {mat_a.size(0), mat_b.size(1)}); + + auto bias_ = bias.value_or(Tensor()); + + if (gemm_impl == ScaledGemmImplementation::TENSORWISE_TENSORWISE || + gemm_impl == ScaledGemmImplementation::ROWWISE_ROWWISE) { + _scaled_mm_out_cpu_emulated( + mat_a, + mat_b, + scale_a[0], + scale_b[0], + bias, + std::nullopt, // scale-result + out_dtype_, + use_fast_accum, + gemm_impl, + out); + } else { + TORCH_CHECK_VALUE(false, "Invalid state - found an implementation, but not really"); + } + + return out; +} + +Tensor _scaled_mm_cpu_v2( + const Tensor& mat_a, + const Tensor& mat_b, + ArrayRef scale_a, + IntArrayRef scale_recipe_a, + IntArrayRef swizzle_a, + ArrayRef scale_b, + IntArrayRef scale_recipe_b, + IntArrayRef swizzle_b, + const std::optional& bias, + const std::optional out_dtype, + IntArrayRef contraction_dim, + bool use_fast_accum) { + const auto out_dtype_ = out_dtype.value_or(mat_a.scalar_type()); + Tensor out = at::empty({0}, mat_a.options().dtype(out_dtype_)); + + return _scaled_mm_cpu_v2_out( + mat_a, + mat_b, + scale_a, + scale_recipe_a, + swizzle_a, + scale_b, + scale_recipe_b, + swizzle_b, + bias, + out_dtype, + contraction_dim, + use_fast_accum, + out); +} + // TODO(vasiliy, future PR): figure out why we need to declare this function, when // other functions that live in ATen/native/*.cpp without declarations // or headers work just fine. diff --git a/aten/src/ATen/native/Sorting.cpp b/aten/src/ATen/native/Sorting.cpp index eb3c2c93ec8a0..9c1d6c4c4e0bc 100644 --- a/aten/src/ATen/native/Sorting.cpp +++ b/aten/src/ATen/native/Sorting.cpp @@ -206,7 +206,7 @@ QUANTILE_INTERPOLATION_MODE get_quantile_interpolation_mode( } void quantile_checks(const Tensor& self, const Tensor& q) { - TORCH_CHECK(self.numel() > 0, "quantile() input tensor must be non-empty"); + TORCH_SYM_CHECK(self.sym_numel().sym_gt(0), "quantile() input tensor must be non-empty"); TORCH_CHECK(q.dim() <= 1, "quantile() q must be a scalar or 1D tensor"); TORCH_CHECK( self.scalar_type() == kFloat || self.scalar_type() == kDouble, @@ -219,26 +219,26 @@ void quantile_checks(const Tensor& self, const Tensor& q) { "quantile() q tensor must be on the same device as the input tensor"); } -std::vector quantile_output_shape( +std::vector quantile_output_shape( const std::optional original_dim, const Tensor& self, const Tensor& q, const bool keepdim, int64_t wrapped_dim) { // Compute output shape: q_size + reduced_size - std::vector out_shape; + std::vector out_shape; if (original_dim && self.dim() > 0) { - out_shape = self.sizes().vec(); + out_shape = self.sym_sizes().vec(); if (keepdim) { out_shape[wrapped_dim] = 1; } else { out_shape.erase(out_shape.begin() + wrapped_dim); } } else if (keepdim) { - out_shape = std::vector(self.dim(), 1); + out_shape = std::vector(self.dim(), 1); } if (q.dim() > 0) { - out_shape.insert(out_shape.begin(), q.numel()); + out_shape.insert(out_shape.begin(), q.sym_numel()); } return out_shape; @@ -252,11 +252,13 @@ Tensor quantile_compute( const QUANTILE_INTERPOLATION_MODE& interpolation, const bool ignore_nan, int64_t wrapped_dim, - std::vector out_shape) { + std::vector out_shape) { // Checks that all q values are between 0 and 1, inclusive // NOTE: this check is only performed when running on the CPU to avoid // synchronizing an accelerator with the CPU - if (self.device().is_cpu()) { + // The check is also skipped when the actual q values are not available yet + // e.g. with symbolic shapes or during export + if (self.device().is_cpu() && !isTensorSubclassLike(q)) { auto all_q_in_range = q.ge(0).logical_and_(q.le(1)).all(); TORCH_CHECK(at::is_scalar_tensor_true(all_q_in_range), "quantile() q values must be in the range [0, 1]"); @@ -275,18 +277,18 @@ Tensor quantile_compute( // Treat q as a 1D tensor for the following computations if (q.dim() == 0) { - out_shape.insert(out_shape.begin(), q.numel()); + out_shape.insert(out_shape.begin(), 1); } // View input as reduced_size + size of dim to reduce - std::vector in_shape(out_shape.size()); + std::vector in_shape(out_shape.size()); std::copy(out_shape.begin() + 1, out_shape.end(), in_shape.begin()); - in_shape[in_shape.size() - 1] = sorted.size(-1); - sorted = sorted.view(in_shape); + in_shape[in_shape.size() - 1] = sorted.sym_size(-1); + sorted = sorted.view_symint(in_shape); // Ensure converting from int64_t to double won't overflow - TORCH_CHECK( - sorted.size(-1) <= std::pow(2, 24), + TORCH_SYM_CHECK( + sorted.sym_size(-1).sym_le(1 << 24), "quantile() input tensor is too large"); // Convert q in [0, 1] to ranks in [0, reduction_size) @@ -308,7 +310,7 @@ Tensor quantile_compute( } else { // For quantile, compute ranks based on reduction size. If there is nan // set rank to last index so the quantile computed will be nan. - int64_t last_index = sorted.size(-1) - 1; + auto last_index = sorted.sym_size(-1) - 1; std::vector tl = at::broadcast_tensors({q * last_index, sorted.isnan().any(-1, true)}); ranks = at::masked_fill(tl[0], tl[1], last_index); @@ -388,7 +390,7 @@ void quantile_out_impl( int64_t wrapped_dim = at::maybe_wrap_dim(original_dim.value_or(0), self.dim()); auto out_shape = quantile_output_shape(original_dim, self, q, keepdim, wrapped_dim); - resize_output(out, out_shape); + resize_output_symint(out, out_shape); auto quantile = quantile_compute( self, q, original_dim, keepdim, interpolation, ignore_nan, wrapped_dim, std::move(out_shape)); diff --git a/aten/src/ATen/native/SpectralOps.cpp b/aten/src/ATen/native/SpectralOps.cpp index 91a0c3ff8cf93..6a8dd6c75769b 100644 --- a/aten/src/ATen/native/SpectralOps.cpp +++ b/aten/src/ATen/native/SpectralOps.cpp @@ -80,9 +80,9 @@ ScalarType promote_type_fft(ScalarType type, bool require_complex, Device device } const bool maybe_support_half = ( - // Only CUDA supports half precision, but since meta tensors don't have a + // CUDA and XPU support half precision, but since meta tensors don't have a // device we err on the side of accepting it - device.is_cuda() || device.is_meta() + device.is_cuda() || device.is_meta() || device.is_xpu() ); if (maybe_support_half) { TORCH_CHECK(type == kHalf || type == kFloat || type == kDouble, "Unsupported dtype ", type); diff --git a/aten/src/ATen/native/TensorAdvancedIndexing.cpp b/aten/src/ATen/native/TensorAdvancedIndexing.cpp index ecb67c9ef3799..4f91a78c0c17c 100644 --- a/aten/src/ATen/native/TensorAdvancedIndexing.cpp +++ b/aten/src/ATen/native/TensorAdvancedIndexing.cpp @@ -988,7 +988,8 @@ Tensor& _index_put_impl_( } } if ((self.device().type() == DeviceType::CUDA || - self.device().type() == DeviceType::XPU) && + self.device().type() == DeviceType::XPU || + self.device().type() == DeviceType::PrivateUse1) && (accumulate || (globalContext().deterministicAlgorithms() && value_.numel() > 1))) { TORCH_CHECK( diff --git a/aten/src/ATen/native/TensorProperties.cpp b/aten/src/ATen/native/TensorProperties.cpp index c15b082f107b2..26cafbb4b2585 100644 --- a/aten/src/ATen/native/TensorProperties.cpp +++ b/aten/src/ATen/native/TensorProperties.cpp @@ -13,10 +13,15 @@ #include #include #include +#include #include +#include +#include #include #include +#include #include +#include #include #include #include @@ -76,6 +81,30 @@ c10::SymInt sym_storage_offset(const Tensor& self) { return self.sym_storage_offset(); } +int64_t numel(const Tensor& self) { + return self.numel(); +} + +int64_t dim(const Tensor& self) { + return self.dim(); +} + +int64_t get_device(const Tensor& self) { + return self.get_device(); +} + +int64_t storage_offset(const Tensor& self) { + return self.storage_offset(); +} + +bool is_contiguous(const Tensor& self) { + return self.is_contiguous(); +} + +bool is_contiguous(const Tensor& self, at::MemoryFormat memory_format) { + return self.is_contiguous(memory_format); +} + int64_t size(const Tensor& self, Dimname dim) { size_t pos_dim = dimname_to_position(self, dim); return self.sizes()[pos_dim]; diff --git a/aten/src/ATen/native/UnaryOps.cpp b/aten/src/ATen/native/UnaryOps.cpp index 8b59cb87730d0..ef534208e61e8 100644 --- a/aten/src/ATen/native/UnaryOps.cpp +++ b/aten/src/ATen/native/UnaryOps.cpp @@ -912,6 +912,10 @@ Tensor& mvlgamma_(Tensor& self, int64_t p) { } Tensor& mvlgamma_out(const Tensor& self, int64_t p, Tensor& result) { + TORCH_CHECK( + self.device() == result.device(), + "Expected tensors to be on the same device, but found ", self.device(), " and ", result.device() + ); auto out = self.mvlgamma(p); TORCH_CHECK( at::can_cast(out.scalar_type(), result.scalar_type()), diff --git a/aten/src/ATen/native/UpSample.h b/aten/src/ATen/native/UpSample.h index cf6727c2207c7..6483ca256414e 100644 --- a/aten/src/ATen/native/UpSample.h +++ b/aten/src/ATen/native/UpSample.h @@ -78,6 +78,7 @@ using _upsampling_bilinear2d_aa = void(*)(const Tensor& output, const Tensor& in using upsampling_trilinear3d = void(*)(const Tensor& output, const Tensor& input, bool align_corners, scale_t scales_d, scale_t scales_h, scale_t scales_w); using upsampling_bicubic2d = void(*)(const Tensor& output, const Tensor& input, bool align_corners, scale_t scales_h, scale_t scales_w); using _upsampling_bicubic2d_aa = void(*)(const Tensor& output, const Tensor& input, bool align_corners, scale_t scales_h, scale_t scales_w); +using _upsampling_lanczos2d_aa = void(*)(const Tensor& output, const Tensor& input, bool align_corners, scale_t scales_h, scale_t scales_w); DECLARE_DISPATCH(upsampling_nearest1d, upsample_nearest1d_kernel) DECLARE_DISPATCH(_upsampling_nearest_exact1d, _upsample_nearest_exact1d_kernel) DECLARE_DISPATCH(upsampling_nearest2d, upsample_nearest2d_kernel) @@ -101,6 +102,8 @@ DECLARE_DISPATCH(upsampling_trilinear3d, upsample_trilinear3d_backward_kernel) DECLARE_DISPATCH(upsampling_bicubic2d, upsample_bicubic2d_kernel) DECLARE_DISPATCH(_upsampling_bicubic2d_aa, _upsample_bicubic2d_aa_kernel) DECLARE_DISPATCH(_upsampling_bicubic2d_aa, _upsample_bicubic2d_aa_backward_kernel) +DECLARE_DISPATCH(_upsampling_lanczos2d_aa, _upsample_lanczos2d_aa_kernel) +DECLARE_DISPATCH(_upsampling_lanczos2d_aa, _upsample_lanczos2d_aa_backward_kernel) [[maybe_unused]] inline std::array upsample_1d_common_check( IntArrayRef input_size, diff --git a/aten/src/ATen/native/UpSampleLanczos2d.cpp b/aten/src/ATen/native/UpSampleLanczos2d.cpp new file mode 100644 index 0000000000000..22ae7d972029f --- /dev/null +++ b/aten/src/ATen/native/UpSampleLanczos2d.cpp @@ -0,0 +1,102 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +#include +#include +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#include +#else +#include +#include +#include +#include +#endif + +namespace at::meta { + +TORCH_META_FUNC(_upsample_lanczos2d_aa) ( + const Tensor& input, IntArrayRef output_size, bool align_corners, std::optional scales_h, std::optional scales_w +) { + auto full_output_size = native::upsample_2d_common_check(input.sizes(), output_size); + + TORCH_CHECK( + input.numel() != 0 || c10::multiply_integers(input.sizes().begin() + 1, input.sizes().end()), + "Non-empty 4D data tensor expected but got a tensor with sizes ", + input.sizes()); + + set_output_raw_strided(0, full_output_size, {}, input.options().memory_format(input.suggest_memory_format())); +} + +TORCH_META_FUNC(_upsample_lanczos2d_aa_backward) ( + const Tensor& grad_output, + IntArrayRef output_size, + IntArrayRef input_size, + bool align_corners, + std::optional scales_h, + std::optional scales_w +) { + auto full_output_size = native::upsample_2d_common_check(input_size, output_size); + + TORCH_CHECK( + grad_output.dim() == 4, + "Expected grad_output to be a tensor of dimension 4 but got: dimension ", grad_output.dim()); + + for (const auto i : c10::irange(4)) { + TORCH_CHECK( + grad_output.size(i) == full_output_size[i], + "Expected grad_output to have the same shape as output;", + " output.size(", i, ") = ", full_output_size[i], + " but got grad_output.size(", i, ") = ", grad_output.size(i)); + } + + set_output_raw_strided(0, input_size, {}, grad_output.options()); +} + +} // namespace at::meta +namespace at::native { + +TORCH_IMPL_FUNC(_upsample_lanczos2d_aa_out_cpu) ( + const Tensor& input, + IntArrayRef output_size, + bool align_corners, + std::optional scales_h, + std::optional scales_w, + const Tensor& output +) { + _upsample_lanczos2d_aa_kernel(kCPU, output, input, align_corners, scales_h, scales_w); +} + +TORCH_IMPL_FUNC(_upsample_lanczos2d_aa_backward_out_cpu) ( + const Tensor& grad_output, + IntArrayRef output_size, + IntArrayRef input_size, + bool align_corners, + std::optional scales_h, + std::optional scales_w, + const Tensor& grad_input +) { + grad_input.zero_(); + _upsample_lanczos2d_aa_backward_kernel(kCPU, grad_input, grad_output, align_corners, scales_h, scales_w); +} + +// vec variant + +using at::native::upsample::compute_output_size; +using at::native::upsample::get_scale_value; + +Tensor _upsample_lanczos2d_aa( + const Tensor& input, + at::OptionalIntArrayRef output_size, + bool align_corners, + std::optional> scale_factors) { + auto osize = compute_output_size(input.sizes(), output_size, scale_factors); + auto scale_h = get_scale_value(scale_factors, 0); + auto scale_w = get_scale_value(scale_factors, 1); + return at::_upsample_lanczos2d_aa(input, osize, align_corners, scale_h, scale_w); +} + +DEFINE_DISPATCH(_upsample_lanczos2d_aa_kernel); +DEFINE_DISPATCH(_upsample_lanczos2d_aa_backward_kernel); + +} // namespace at::native diff --git a/aten/src/ATen/native/cpu/Loops.h b/aten/src/ATen/native/cpu/Loops.h index 05e7f93d2f364..adadb636e7f9b 100644 --- a/aten/src/ATen/native/cpu/Loops.h +++ b/aten/src/ATen/native/cpu/Loops.h @@ -301,13 +301,15 @@ VectorizedLoop2d make_vectorized_loop2d( } template -void cpu_kernel(TensorIteratorBase& iter, func_t&& op, int64_t grain_size = at::internal::GRAIN_SIZE) { +void cpu_kernel(TensorIteratorBase& iter, func_t&& op, int64_t grain_size = at::internal::GRAIN_SIZE, bool check_dynamic_casting = true) { using traits = function_traits; // this could be extended to work with void return types TORCH_INTERNAL_ASSERT(iter.ninputs() == traits::arity); TORCH_INTERNAL_ASSERT(iter.noutputs() == 1); // dynamic casting not currently supported on CPU - TORCH_INTERNAL_ASSERT(!needs_dynamic_casting::check(iter)); + if (check_dynamic_casting) { + TORCH_INTERNAL_ASSERT(!needs_dynamic_casting::check(iter)); + } iter.for_each([&](char** data, const int64_t* strides, int64_t n) { // basic loop can handle 1d slices with arbitrary strides, and 1d slices is all that @@ -317,6 +319,11 @@ void cpu_kernel(TensorIteratorBase& iter, func_t&& op, int64_t grain_size = at:: iter.cast_outputs(); } +template +void cpu_kernel_opaque(TensorIteratorBase& iter, func_t&& op, int64_t grain_size = at::internal::GRAIN_SIZE) { + return cpu_kernel(iter, op, grain_size, false); +} + // This function helps write elementwise kernels that requires multiple outputs. // It follows the similar structure of cpu_kernel. // Instead of `basic_loop` function, a new `multiple_outputs_loop` function is diff --git a/aten/src/ATen/native/cpu/UpSampleKernel.cpp b/aten/src/ATen/native/cpu/UpSampleKernel.cpp index 39fe91ccc06e4..6274f22ecad7c 100644 --- a/aten/src/ATen/native/cpu/UpSampleKernel.cpp +++ b/aten/src/ATen/native/cpu/UpSampleKernel.cpp @@ -24,21 +24,54 @@ namespace { using scale_t = std::vector>; +// Naming conventions used in this file: +// +// - "non_separable": All spatial dimensions are interpolated in a single +// TensorIterator pass, using the recursive InterpolateNonSeparable struct. +// Non-separable is always multi-dimensional (Nd). +// +// - "separable": The multi-dimensional interpolation is decomposed into a +// sequence of 1d passes (one per spatial dimension). The entry point is Nd +// (it loops over dims), and inner functions are 1d. +// +// - "1d": Refers to processing a single spatial dimension within the +// separable approach. Not to be confused with 1d interpolation (e.g. +// linear); it means one dimension of a multi-dimensional separable +// decomposition. +// +// - "Nd": Templated on out_ndims (1, 2, or 3) for the separable approach. + + +// ---- Non-separable interpolation ---- +// +// Used by: nearest, linear, bilinear (float), cubic (float), trilinear. +// Processes all spatial dims in a single TensorIterator pass via the recursive +// InterpolateNonSeparable struct. +// +// Call chain: +// upsample_non_separable_Nd_kernel_impl +// -> upsample_non_separable +// -> basic_loop_non_separable +// -> interpolate_non_separable +// -> InterpolateNonSeparable (recursive struct) +// +// Helper structs and methods for upsample_non_separable +// // Interpolation structure to compute output value in n-dimensional case. // - recursively compute interpolated output for each dimension // - we rely a lot on compiler's code optimization such that implemented operations // can be automatically factorized and vectorized using SSE and AVX2 template -struct Interpolate { +struct InterpolateNonSeparable { static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) { index_t ids = *(index_t*)&data[0][i * strides[0]]; opmath_t wts = *(scalar_t*)&data[1][i * strides[1]]; - opmath_t t = Interpolate::eval(src + ids, &data[2 * interp_size], &strides[2 * interp_size], i); + opmath_t t = InterpolateNonSeparable::eval(src + ids, &data[2 * interp_size], &strides[2 * interp_size], i); opmath_t output = t * wts; for (const auto j : c10::irange(1, interp_size)) { ids = *(index_t*)&data[2 * j + 0][i * strides[2 * j + 0]]; wts = *(scalar_t*)&data[2 * j + 1][i * strides[2 * j + 1]]; - t = Interpolate::eval(src + ids, &data[2 * interp_size], &strides[2 * interp_size], i); + t = InterpolateNonSeparable::eval(src + ids, &data[2 * interp_size], &strides[2 * interp_size], i); output += t * wts; } return output; @@ -46,7 +79,7 @@ struct Interpolate { }; template -struct Interpolate<1, scalar_t, opmath_t, index_t, interp_size> { +struct InterpolateNonSeparable<1, scalar_t, opmath_t, index_t, interp_size> { static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) { index_t ids = *(index_t*)&data[0][i * strides[0]]; opmath_t wts = *(scalar_t*)&data[1][i * strides[1]]; @@ -63,15 +96,15 @@ struct Interpolate<1, scalar_t, opmath_t, index_t, interp_size> { }; template -struct Interpolate { +struct InterpolateNonSeparable { static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) { index_t ids = *(index_t*)&data[0][i * strides[0]]; - return Interpolate::eval(src + ids, &data[2], &strides[2], i); + return InterpolateNonSeparable::eval(src + ids, &data[2], &strides[2], i); } }; template -struct Interpolate<1, scalar_t, opmath_t, index_t, 1> { +struct InterpolateNonSeparable<1, scalar_t, opmath_t, index_t, 1> { static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) { index_t ids = *(index_t*)&data[0][i * strides[0]]; return *(scalar_t *)&src[ids]; @@ -81,25 +114,25 @@ struct Interpolate<1, scalar_t, opmath_t, index_t, 1> { // There is an unexpected 2x slowdown for upsample_trilinear3d channels_first // for both 1 and 6 threads. We have to specialize this case as below: // Once the issue is fixed we can keep generic implementation and remove: -// struct Interpolate and -// struct Interpolate<1, scalar_t, index_t, 2> +// struct InterpolateNonSeparable and +// struct InterpolateNonSeparable<1, scalar_t, index_t, 2> template -struct Interpolate { +struct InterpolateNonSeparable { static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) { index_t i0 = *(index_t*)&data[0][i * strides[0]]; index_t i1 = *(index_t*)&data[2][i * strides[2]]; opmath_t w0 = *(scalar_t *)&data[1][i * strides[1]]; opmath_t w1 = *(scalar_t *)&data[3][i * strides[3]]; - opmath_t t0 = Interpolate::eval(src + i0, &data[4], &strides[4], i); - opmath_t t1 = Interpolate::eval(src + i1, &data[4], &strides[4], i); + opmath_t t0 = InterpolateNonSeparable::eval(src + i0, &data[4], &strides[4], i); + opmath_t t1 = InterpolateNonSeparable::eval(src + i1, &data[4], &strides[4], i); return t0 * w0 + t1 * w1; } }; template -struct Interpolate<1, scalar_t, opmath_t, index_t, 2> { +struct InterpolateNonSeparable<1, scalar_t, opmath_t, index_t, 2> { static inline opmath_t eval(char* src, char** data, const int64_t* strides, int64_t i) { index_t i0 = *(index_t*)&data[0][i * strides[0]]; index_t i1 = *(index_t*)&data[2][i * strides[2]]; @@ -112,13 +145,28 @@ struct Interpolate<1, scalar_t, opmath_t, index_t, 2> { }; template -inline scalar_t interpolate(char* src, char** data, const int64_t* strides, int64_t i) { +inline scalar_t interpolate_non_separable(char* src, char** data, const int64_t* strides, int64_t i) { using opmath_t = at::opmath_type; - return Interpolate::eval(src, data, strides, i); + return InterpolateNonSeparable::eval(src, data, strides, i); } +// ---- Separable interpolation ---- +// +// Used by: bilinear/bicubic with antialias=True, and bilinear/bicubic uint8 +// (as fallback when AVX isn't supported). +// Processes one spatial dimension at a time. The outer loop over dimensions +// is in upsample_separable_Nd_kernel_impl. +// +// Call chain: +// upsample_separable_Nd_kernel_impl (loops over dims) +// -> upsample_separable_1d +// -> basic_loop_separable_1d_horizontal (for last spatial dim, i.e. W) +// -> interpolate_separable_1d +// -> basic_loop_separable_1d_vertical (for other spatial dims, e.g. H, D) +// -> interpolate_separable_1d_zero_strides + template -inline scalar_t interpolate_aa_single_dim_zero_strides( +inline scalar_t interpolate_separable_1d_zero_strides( char* src, char** data, const index_t ids_stride) { @@ -142,7 +190,7 @@ inline scalar_t interpolate_aa_single_dim_zero_strides( } template -inline scalar_t interpolate_aa_single_dim( +inline scalar_t interpolate_separable_1d( char* src, char** data, const int64_t* strides, @@ -215,7 +263,7 @@ inline bool is_contiguous_stride(const int64_t* strides) { // strides=(0, 0, 0, 0, 0, 0, 0, 0, ..., 0, 0, 0, 0) // // Using these methods we can hint the compiler to factorize constant indices and weights -// in cpu_upsample_linear method +// in upsample_non_separable template struct CheckAlmostAllZeroStrides { static inline bool eval(const int64_t* strides) { @@ -247,17 +295,17 @@ inline bool check_almost_all_zero_stride(const int64_t* strides) { // Helper method to compute interpolation for nearest, linear, cubic modes template -inline void basic_loop(char** data, const int64_t* strides, int64_t n) { +inline void basic_loop_non_separable(char** data, const int64_t* strides, int64_t n) { char* dst = data[0]; char* src = data[1]; for (const auto i : c10::irange(n)) { - *(scalar_t*)&dst[i * strides[0]] = interpolate( + *(scalar_t*)&dst[i * strides[0]] = interpolate_non_separable( src + i * strides[1], &data[2], &strides[2], i); } } template -inline void basic_loop_aa_vertical( +inline void basic_loop_separable_1d_vertical( char** data, const int64_t* strides, int64_t n, @@ -269,13 +317,13 @@ inline void basic_loop_aa_vertical( for (const auto i : c10::irange(n)) { *(scalar_t*)&dst[i * strides[0]] = - interpolate_aa_single_dim_zero_strides( + interpolate_separable_1d_zero_strides( src + i * strides[1], &data[2], ids_stride); } } template <> -inline void basic_loop_aa_vertical( +inline void basic_loop_separable_1d_vertical( char** data, const int64_t* strides, int64_t n, @@ -313,7 +361,7 @@ inline void basic_loop_aa_vertical( } template -inline void basic_loop_aa_horizontal( +inline void basic_loop_separable_1d_horizontal( char** data, const int64_t* strides, int64_t n, @@ -326,20 +374,20 @@ inline void basic_loop_aa_horizontal( if (strides[1] == 0) { for (const auto i : c10::irange(n)) { *(scalar_t*)&dst[i * strides[0]] = - interpolate_aa_single_dim( + interpolate_separable_1d( src, &data[2], &strides[2], i, ids_stride); } } else { for (const auto i : c10::irange(n)) { *(scalar_t*)&dst[i * strides[0]] = - interpolate_aa_single_dim( + interpolate_separable_1d( src + i * strides[1], &data[2], &strides[2], i, ids_stride); } } } template <> -inline void basic_loop_aa_horizontal( +inline void basic_loop_separable_1d_horizontal( char** data, const int64_t* strides, int64_t n, @@ -390,10 +438,10 @@ inline void basic_loop_aa_horizontal( // output_DN[a] = interpolate(input_DN[a], w_DN[a], input_DN[a+1], w_DN[a+1], ...) // and i - dimension index and a - linear index for spatial coordinates // -// The recursive call is implemented with InterpLinear struct using template for +// The recursive call is implemented with the InterpolateNonSeparable struct using template for // the loop unrolling on compile time. template -void cpu_upsample_generic(at::TensorIterator& iter) +void upsample_non_separable(at::TensorIterator& iter) { auto loop = [&](char** data, const int64_t* strides, int64_t n) { // special-cases to let the compiler apply compile-time input-specific optimizations @@ -401,21 +449,21 @@ void cpu_upsample_generic(at::TensorIterator& iter) // NOLINTNEXTLINE(bugprone-branch-clone) check_almost_all_zero_stride(&strides[2]))) { // contiguous channels-first case - basic_loop(data, strides, n); + basic_loop_non_separable(data, strides, n); } else if ((strides[0] == sizeof(scalar_t) && (strides[1] == sizeof(scalar_t)) && check_almost_all_zero_stride(&strides[2]))) { // contiguous channels-last case - basic_loop(data, strides, n); + basic_loop_non_separable(data, strides, n); } else { // fallback - basic_loop(data, strides, n); + basic_loop_non_separable(data, strides, n); } }; iter.for_each(loop); } template -void cpu_upsample_nearest_channels_last( +void upsample_nearest_channels_last( const Tensor& output_, const Tensor& input_, const scale_type& scales) { @@ -520,7 +568,7 @@ inline VecType interpolate(const scalar_t* t, accscalar_t w, Args... a } template -void cpu_upsample_linear_channels_last( +void upsample_linear_channels_last( const Tensor& output_, const Tensor& input_, bool align_corners, @@ -680,7 +728,7 @@ void cpu_upsample_linear_channels_last( } } -// Helper structs to use with upsample_generic_Nd_kernel_impl +// Helper structs to use with upsample_non_separable_Nd_kernel_impl struct HelperInterpBase { static inline void init_indices_weights( @@ -746,7 +794,7 @@ struct HelperInterpBase { // for interpolation with antialiasing=false mode. It returns the maximal weights value. // This function is templated with scalar_t for type of scale and weights but is only used for // bilinear/bicubic modes on uint8 input and antialiasing=false (in this case scalar_t is double). - // For float input types we are using upsample_generic_Nd_kernel_impl and compute_indices_weights methods + // For float input types we are using upsample_non_separable_Nd_kernel_impl and compute_indices_weights methods template static inline scalar_t _compute_indices_min_size_weights( const int64_t i, const int64_t input_size, const scalar_t scale, @@ -906,7 +954,7 @@ struct HelperInterpBase { weights as double, but then convert them to int16 via some conversion logic detailed below. This allows us to compute all interpolation operation (sum of multiplications) as ints instead of floats. The result is converted back into - uint8 in basic_loop_aa_horizontal (and vertical) + uint8 in basic_loop_separable_1d_horizontal (and vertical) In essence the idea is to avoid a multiplication between a float (the weight) and an int (the pixel value) and instead run a multiplication between @@ -1368,6 +1416,89 @@ struct HelperInterpCubic : public HelperInterpBase { }; +struct HelperInterpLanczos : public HelperInterpBase { + + static constexpr int interp_size = 6; + + // Taken from + // https://github.com/python-pillow/Pillow/blob/8004234d879254cc354935ad42fbb51b1700925e/src/libImaging/Resample.c#L64-L86 + template + static inline scalar_t sinc_filter(scalar_t x) { + if (x == 0.0) { + return 1.0; + } + x *= c10::pi; + return std::sin(x) / x; + } + + template + static inline scalar_t aa_filter(scalar_t x) { + // Lanczos-3 filter: sinc(x) * sinc(x/3) for |x| < 3 + x = std::abs(x); + if (x < 3.0) { + return sinc_filter(x) * sinc_filter(x / 3.0); + } + return 0.0; + } + + static inline std::vector compute_index_ranges_weights( + at::ScalarType scalar_type, + int64_t input_size, + int64_t output_size, + int64_t stride, + int64_t ndims, + int64_t reshape_dim, + bool align_corners, + const std::optional& opt_scale, + bool antialias + ) { + + std::vector indices_weights; + AT_DISPATCH_FLOATING_TYPES( + scalar_type, "compute_index_ranges_weights", [&] { + + scalar_t scale = area_pixel_compute_scale( + input_size, output_size, align_corners, opt_scale); + + auto interp_size = HelperInterpLanczos::interp_size; + + indices_weights = std::get<0>(HelperInterpLanczos::_compute_index_ranges_weights( + input_size, + output_size, + stride, + ndims, + reshape_dim, + scale, + interp_size, + &HelperInterpLanczos::aa_filter, + /*antialias=*/antialias, + /*align_corners=*/align_corners)); + } + ); + return indices_weights; + } + + static inline std::tuple, int, unsigned int> compute_index_ranges_int16_weights( + int64_t input_size, + int64_t output_size, + int64_t stride, + int64_t ndims, + int64_t reshape_dim, + bool align_corners, + const std::optional& opt_scale, + bool antialias, + bool align_i32=false + ) { + + auto interp_size = HelperInterpLanczos::interp_size; + auto fn = HelperInterpLanczos::aa_filter; + return HelperInterpLanczos::_compute_index_ranges_int16_weights( + input_size, output_size, stride, ndims, reshape_dim, + align_corners, opt_scale, interp_size, fn, antialias, align_i32); + } + +}; + // Generic upsampling interpolation kernel for N-d case. // Input is assumed to be like NCHW, NCL, NCKHW - interpolated spatial dimension // are those from the end up to batch size N and number of channels C. @@ -1377,7 +1508,7 @@ struct HelperInterpCubic : public HelperInterpBase { // - scale_type is template type for scales, typically std::optional // - template class F is one of the above structs to compute indices and weights template -void upsample_generic_Nd_kernel_impl( +void upsample_non_separable_Nd_kernel_impl( const Tensor& output, const Tensor& input, bool align_corners, @@ -1437,43 +1568,27 @@ void upsample_generic_Nd_kernel_impl( if (interp_size > 1) { // Nearest also supports uint8 tensor, so need to handle it separately + // Dispatch name should be "upsample_non_separable" but we keep the old + // name for internal BC. AT_DISPATCH_FLOATING_TYPES_AND2( kBFloat16, kHalf, iter.dtype(), "upsample_generic_Nd", [&] { // MSVC can not catch constexpr int interp_size here constexpr int mode = F::interp_size; - cpu_upsample_generic(iter); + upsample_non_separable(iter); }); } else { + // Dispatch name should be "upsample_non_separable" but we keep the old + // name for internal BC. AT_DISPATCH_FLOATING_TYPES_AND3(kByte, kBFloat16, kHalf, iter.dtype(), "upsample_generic_Nd", [&] { constexpr int mode = F::interp_size; - cpu_upsample_generic(iter); + upsample_non_separable(iter); }); } } -template -void cpu_upsample_generic_aa(at::TensorIterator& iter, unsigned int weights_precision) { - - auto loop = [&](char** data, const int64_t* strides, int64_t n) { - if constexpr (is_horizontal) { - - // Strides are : X 0 | 8 8 8 0 8 (Channels first) - // Strides are : X X | 0 0 0 0 0 (Channels last) - basic_loop_aa_horizontal(data, strides, n, weights_precision); - } else { - // Strides are : X Y | 0 0 0 0 0 (Channels first) - // Strides are : X X | 0 0 0 0 0 (Channels last) - // upsampling data between contiguous dimensions (aka vertical resampling) - basic_loop_aa_vertical(data, strides, n, weights_precision); - } - }; - - iter.for_each(loop); -} - template -void _separable_upsample_generic_Nd_kernel_impl_single_dim( +void upsample_separable_1d( const Tensor& output, const Tensor& input, int interp_dim, @@ -1502,8 +1617,8 @@ void _separable_upsample_generic_Nd_kernel_impl_single_dim( unsigned int weights_precision = 0; if (input_scalar_type == at::kByte) { - // This is a special branch to provide uint8 dtype support for bilinear and bicubic modes only - TORCH_INTERNAL_ASSERT(F::interp_size == 2 || F::interp_size == 4); + // This is a special branch to provide uint8 dtype support for bilinear, bicubic and lanczos modes only + TORCH_INTERNAL_ASSERT(F::interp_size == 2 || F::interp_size == 4 || F::interp_size == 6); int unused = 0; std::tie(indices_weights, unused, weights_precision) = F::compute_index_ranges_int16_weights( @@ -1533,9 +1648,22 @@ void _separable_upsample_generic_Nd_kernel_impl_single_dim( auto iter = config.build(); + // Dispatch name should be "upsample_separable_1d" but we keep the old + // name for internal BC. AT_DISPATCH_FLOATING_TYPES_AND( at::ScalarType::Byte, iter.dtype(), "upsample_generic_Nd_aa", [&] { - cpu_upsample_generic_aa(iter, weights_precision); + auto loop = [&](char** data, const int64_t* strides, int64_t n) { + if constexpr (is_horizontal) { + // Strides are : X 0 | 8 8 8 0 8 (Channels first) + // Strides are : X X | 0 0 0 0 0 (Channels last) + basic_loop_separable_1d_horizontal(data, strides, n, weights_precision); + } else { + // Strides are : X Y | 0 0 0 0 0 (Channels first) + // Strides are : X X | 0 0 0 0 0 (Channels last) + basic_loop_separable_1d_vertical(data, strides, n, weights_precision); + } + }; + iter.for_each(loop); }); } @@ -1544,7 +1672,7 @@ void _separable_upsample_generic_Nd_kernel_impl_single_dim( // (dtype == uint8 and mode in ("bilinear", "bicubic")): this is used as // fallback in these settings when AVX isn't supported. template -void separable_upsample_generic_Nd_kernel_impl( +void upsample_separable_Nd_kernel_impl( const Tensor& output, const Tensor& input, bool align_corners, @@ -1563,29 +1691,29 @@ void separable_upsample_generic_Nd_kernel_impl( at::Tensor temp_output, temp_input = input; int interp_dim = 0; - // Precompute the number of single dim resize method invocations + // Precompute the number of 1d resize ops // to avoid copying temporary buffer to output - int num_single_dim_ops = 0; + int num_1d_ops = 0; for (const auto i : c10::irange(out_ndims)) { interp_dim = 2 + out_ndims - 1 - i; if (output_shape[interp_dim] != input_shape[interp_dim]) { - num_single_dim_ops += 1; + num_1d_ops += 1; } } - // upsampling data within the contiguous dimension (aka horizontal resampling) + // Horizontal resampling (last spatial dim, i.e. W) interp_dim = 2 + out_ndims - 1; if (output_shape[interp_dim] != input_shape[interp_dim]) { - num_single_dim_ops -= 1; - if (num_single_dim_ops > 0) { + num_1d_ops -= 1; + if (num_1d_ops > 0) { temp_oshape[interp_dim] = output_shape[interp_dim]; temp_output = at::empty(temp_oshape, input.options()); } else { temp_output = output; } - _separable_upsample_generic_Nd_kernel_impl_single_dim< + upsample_separable_1d< out_ndims, scale_t, F, @@ -1594,20 +1722,20 @@ void separable_upsample_generic_Nd_kernel_impl( temp_input = temp_output; } - // upsampling data between contiguous dimensions (aka vertical resampling) + // Vertical resampling (remaining spatial dims, e.g. H, D) for (const auto i : c10::irange(1, out_ndims)) { interp_dim = 2 + out_ndims - 1 - i; if (output_shape[interp_dim] != input_shape[interp_dim]) { - num_single_dim_ops -= 1; - if (num_single_dim_ops > 0) { + num_1d_ops -= 1; + if (num_1d_ops > 0) { temp_oshape[interp_dim] = output_shape[interp_dim]; temp_output = at::empty(temp_oshape, input.options()); } else { temp_output = output; } - _separable_upsample_generic_Nd_kernel_impl_single_dim< + upsample_separable_1d< out_ndims, scale_t, F, @@ -1622,7 +1750,7 @@ void upsample_nearest1d_kernel_impl( const Tensor& output, const Tensor& input, std::optional scales_w) { - upsample_generic_Nd_kernel_impl<1, scale_t, HelperInterpNearest>( + upsample_non_separable_Nd_kernel_impl<1, scale_t, HelperInterpNearest>( output, input, false, {scales_w}); } @@ -1630,26 +1758,28 @@ void _upsample_nearest_exact1d_kernel_impl( const Tensor& output, const Tensor& input, std::optional scales_w) { - upsample_generic_Nd_kernel_impl<1, scale_t, HelperInterpNearestExact>( + upsample_non_separable_Nd_kernel_impl<1, scale_t, HelperInterpNearestExact>( output, input, false, {scales_w}); } -int _use_vectorized_kernel_cond_2d( +int _use_channels_last_kernel_2d( const Tensor& output, const Tensor& input) { - // This condition is used to know whether we should dispatch to a vectorized - // kernel, or to the more general upsample_generic_Nd_kernel_impl(). For now, - // the vectorized kernels are only optimized for channels_last and when C >= 4 - // (shape = NCHW). For a very wide range of use-cases (typically image or mask - // resizing where we have C < 4), using upsample_generic_Nd_kernel_impl() is - // actually faster. On top of that, benchmarks showed that this also depends on - // the *output* size (output_H + output_W), for both upsampling and - // downsampling. The current 128 threshold was determined through benchmarks. + // This condition is used to know whether we should dispatch to a + // channels-last-optimized kernel, or to the more general + // upsample_non_separable_Nd_kernel_impl(). For now, the channels-last kernels + // are only optimized for channels_last and when C >= 4 (shape = NCHW). + // For a very wide range of use-cases (typically image or mask resizing + // where we have C < 4), using upsample_non_separable_Nd_kernel_impl() is + // actually faster. On top of that, benchmarks showed that this also + // depends on the *output* size (output_H + output_W), for both + // upsampling and downsampling. The current 128 threshold was determined + // through benchmarks. return ((input.is_contiguous(at::MemoryFormat::ChannelsLast)) && (input.size(1) > 3)) || ((output.size(-2) + output.size(-1)) <= 128); } -int _use_vectorized_kernel_cond_3d( - // Similar to _use_vectorized_kernel_cond_2d() but for 3d resampling (e.g. videos) +int _use_channels_last_kernel_3d( + // Similar to _use_channels_last_kernel_2d() but for 3d resampling (e.g. videos) // Note that unlike the 2d case, this is not subject to small output size // overhead - hence the absence of the 128 threshold in the condition. const Tensor& output, @@ -1663,13 +1793,13 @@ void upsample_nearest2d_kernel_impl( const Tensor& input, std::optional scales_h, std::optional scales_w) { - if (_use_vectorized_kernel_cond_2d(output, input)) { + if (_use_channels_last_kernel_2d(output, input)) { AT_DISPATCH_FLOATING_TYPES_AND3(kByte, kBFloat16, kHalf, input.scalar_type(), "upsample_nearest2d_channels_last", [&] { - cpu_upsample_nearest_channels_last(output, input, {scales_h, scales_w}); + upsample_nearest_channels_last(output, input, {scales_h, scales_w}); }); } else { - upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpNearest>( + upsample_non_separable_Nd_kernel_impl<2, scale_t, HelperInterpNearest>( output, input, false, {scales_h, scales_w}); } } @@ -1679,12 +1809,12 @@ void _upsample_nearest_exact2d_kernel_impl( const Tensor& input, std::optional scales_h, std::optional scales_w) { - if (_use_vectorized_kernel_cond_2d(output, input)) { + if (_use_channels_last_kernel_2d(output, input)) { AT_DISPATCH_FLOATING_TYPES_AND3(kByte, kBFloat16, kHalf, input.scalar_type(), "upsample_nearest2d_channels_last", [&] { - cpu_upsample_nearest_channels_last(output, input, {scales_h, scales_w}); + upsample_nearest_channels_last(output, input, {scales_h, scales_w}); }); } else { - upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpNearestExact>( + upsample_non_separable_Nd_kernel_impl<2, scale_t, HelperInterpNearestExact>( output, input, false, {scales_h, scales_w}); } } @@ -1695,13 +1825,13 @@ void upsample_nearest3d_kernel_impl( std::optional scales_d, std::optional scales_h, std::optional scales_w) { - if (_use_vectorized_kernel_cond_3d(output, input)) { + if (_use_channels_last_kernel_3d(output, input)) { AT_DISPATCH_FLOATING_TYPES_AND3(kByte, kBFloat16, kHalf, input.scalar_type(), "upsample_nearest3d_channels_last", [&] { - cpu_upsample_nearest_channels_last(output, input, {scales_d, scales_h, scales_w}); + upsample_nearest_channels_last(output, input, {scales_d, scales_h, scales_w}); }); } else { - upsample_generic_Nd_kernel_impl<3, scale_t, HelperInterpNearest>( + upsample_non_separable_Nd_kernel_impl<3, scale_t, HelperInterpNearest>( output, input, false, {scales_d, scales_h, scales_w}); } } @@ -1712,12 +1842,12 @@ void _upsample_nearest_exact3d_kernel_impl( std::optional scales_d, std::optional scales_h, std::optional scales_w) { - if (_use_vectorized_kernel_cond_3d(output, input)) { + if (_use_channels_last_kernel_3d(output, input)) { AT_DISPATCH_FLOATING_TYPES_AND3(kByte, kBFloat16, kHalf, input.scalar_type(), "upsample_nearest3d_channels_last", [&] { - cpu_upsample_nearest_channels_last(output, input, {scales_d, scales_h, scales_w}); + upsample_nearest_channels_last(output, input, {scales_d, scales_h, scales_w}); }); } else { - upsample_generic_Nd_kernel_impl<3, scale_t, HelperInterpNearestExact>( + upsample_non_separable_Nd_kernel_impl<3, scale_t, HelperInterpNearestExact>( output, input, false, {scales_d, scales_h, scales_w}); } } @@ -1727,7 +1857,7 @@ void upsample_linear1d_kernel_impl( const Tensor& input, bool align_corners, std::optional scales_w) { - upsample_generic_Nd_kernel_impl<1, scale_t, HelperInterpLinear>( + upsample_non_separable_Nd_kernel_impl<1, scale_t, HelperInterpLinear>( output, input, align_corners, {scales_w}); } @@ -1739,17 +1869,17 @@ void upsample_bilinear2d_kernel_impl_float( std::optional scales_h, std::optional scales_w) { - // See note above about _use_vectorized_kernel_cond_2d(output, input). The extra cond is present + // See note above about _use_channels_last_kernel_2d(output, input). The extra cond is present // because benchmarks showed that with only 1 thread, images (C == 3) were - // slightly faster with the vectorized kernel than with the generic one. + // slightly faster with the channels-last kernel than with the generic one. // That's not the case for masks though (C == 1), which strongly benefit from // using the generic kernel. - if ((_use_vectorized_kernel_cond_2d(output, input)) || (at::get_num_threads() == 1 && input.size(1) == 3)) { + if ((_use_channels_last_kernel_2d(output, input)) || (at::get_num_threads() == 1 && input.size(1) == 3)) { AT_DISPATCH_FLOATING_TYPES_AND2(kBFloat16, kHalf, input.scalar_type(), "upsample_bilinear2d_channels_last", [&] { - cpu_upsample_linear_channels_last(output, input, align_corners, {scales_h, scales_w}); + upsample_linear_channels_last(output, input, align_corners, {scales_h, scales_w}); }); } else { - upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpLinear>( + upsample_non_separable_Nd_kernel_impl<2, scale_t, HelperInterpLinear>( output, input, align_corners, {scales_h, scales_w}); } } @@ -1777,7 +1907,7 @@ void upsample_bilinear2d_kernel_impl( /*antialias=*/false); } #endif // CPU_CAPABILITY_AVX2 - return separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpLinear>( + return upsample_separable_Nd_kernel_impl<2, scale_t, HelperInterpLinear>( output, input, align_corners, {scales_h, scales_w}, /*antialias=*/false); } @@ -1808,7 +1938,7 @@ void upsample_bilinear2d_aa_kernel_impl( } #endif // CPU_CAPABILITY_AVX2 } - return separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpLinear>( + return upsample_separable_Nd_kernel_impl<2, scale_t, HelperInterpLinear>( output, input, align_corners, {scales_h, scales_w}, /*antialias=*/true); } @@ -1820,12 +1950,12 @@ void upsample_trilinear3d_kernel_impl( std::optional scales_d, std::optional scales_h, std::optional scales_w) { - if ((_use_vectorized_kernel_cond_3d(output, input))) { + if ((_use_channels_last_kernel_3d(output, input))) { AT_DISPATCH_FLOATING_TYPES_AND2(kBFloat16, kHalf, input.scalar_type(), "upsample_trilinear3d_channels_last", [&] { - cpu_upsample_linear_channels_last(output, input, align_corners, {scales_d, scales_h, scales_w}); + upsample_linear_channels_last(output, input, align_corners, {scales_d, scales_h, scales_w}); }); } else { - upsample_generic_Nd_kernel_impl<3, scale_t, HelperInterpLinear>( + upsample_non_separable_Nd_kernel_impl<3, scale_t, HelperInterpLinear>( output, input, align_corners, {scales_d, scales_h, scales_w}); } } @@ -1853,11 +1983,11 @@ void upsample_bicubic2d_kernel_impl( /*antialias=*/false); } #endif // CPU_CAPABILITY_AVX2 - return separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpCubic>( + return upsample_separable_Nd_kernel_impl<2, scale_t, HelperInterpCubic>( output, input, align_corners, {scales_h, scales_w}, /*antialias=*/false); } - return upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpCubic>( + return upsample_non_separable_Nd_kernel_impl<2, scale_t, HelperInterpCubic>( output, input, align_corners, {scales_h, scales_w}); } @@ -1885,7 +2015,36 @@ void upsample_bicubic2d_aa_kernel_impl( } #endif // CPU_CAPABILITY_AVX2 } - return separable_upsample_generic_Nd_kernel_impl<2, scale_t, HelperInterpCubic>( + return upsample_separable_Nd_kernel_impl<2, scale_t, HelperInterpCubic>( + output, input, align_corners, {scales_h, scales_w}, + /*antialias=*/true); +} + +void upsample_lanczos2d_aa_kernel_impl( + const Tensor& output, + const Tensor& input, + bool align_corners, + std::optional scales_h, + std::optional scales_w) { + + if (input.dtype() == at::kByte) { + #ifdef CPU_CAPABILITY_AVX2 + if (input.size(1) <= 4) { + return upsample_avx_bilinear_bicubic_uint8( + input, output, align_corners, {scales_h, scales_w}, + /*antialias=*/true); + } + #elif defined(__aarch64__) + if (input.size(1) == 3 + && input.is_contiguous(at::MemoryFormat::ChannelsLast) + && output.is_contiguous(at::MemoryFormat::ChannelsLast)) { + return upsample_neon_bilinear_bicubic_uint8( + input, output, align_corners, {scales_h, scales_w}, + /*antialias=*/true); + } + #endif // CPU_CAPABILITY_AVX2 + } + return upsample_separable_Nd_kernel_impl<2, scale_t, HelperInterpLanczos>( output, input, align_corners, {scales_h, scales_w}, /*antialias=*/true); } @@ -1894,7 +2053,7 @@ template < typename scalar_t, typename scale_type, class F> -void cpu_upsample_genNd_backward_aa( +void upsample_separable_Nd_backward_aa( const Tensor& grad_input_, const Tensor& grad_output_, bool align_corners, @@ -2012,7 +2171,7 @@ void upsample_bilinear2d_aa_backward_kernel_impl( std::optional scales_w) { AT_DISPATCH_FLOATING_TYPES( grad_output.scalar_type(), "upsample_bilinear2d_aa_backward_cpu", [&] { - cpu_upsample_genNd_backward_aa( + upsample_separable_Nd_backward_aa( grad_input, grad_output, align_corners, {scales_h, scales_w}); }); } @@ -2025,7 +2184,20 @@ void upsample_bicubic2d_aa_backward_kernel_impl( std::optional scales_w) { AT_DISPATCH_FLOATING_TYPES( grad_output.scalar_type(), "upsample_bicubic2d_aa_backward_cpu", [&] { - cpu_upsample_genNd_backward_aa( + upsample_separable_Nd_backward_aa( + grad_input, grad_output, align_corners, {scales_h, scales_w}); + }); +} + +void upsample_lanczos2d_aa_backward_kernel_impl( + const Tensor& grad_input, + const Tensor& grad_output, + bool align_corners, + std::optional scales_h, + std::optional scales_w) { + AT_DISPATCH_FLOATING_TYPES( + grad_output.scalar_type(), "upsample_lanczos2d_aa_backward_cpu", [&] { + upsample_separable_Nd_backward_aa( grad_input, grad_output, align_corners, {scales_h, scales_w}); }); } @@ -2048,4 +2220,7 @@ REGISTER_DISPATCH(upsample_trilinear3d_kernel, &upsample_trilinear3d_kernel_impl REGISTER_DISPATCH(upsample_bicubic2d_kernel, &upsample_bicubic2d_kernel_impl) REGISTER_DISPATCH(_upsample_bicubic2d_aa_kernel, &upsample_bicubic2d_aa_kernel_impl) REGISTER_DISPATCH(_upsample_bicubic2d_aa_backward_kernel, &upsample_bicubic2d_aa_backward_kernel_impl) + +REGISTER_DISPATCH(_upsample_lanczos2d_aa_kernel, &upsample_lanczos2d_aa_kernel_impl) +REGISTER_DISPATCH(_upsample_lanczos2d_aa_backward_kernel, &upsample_lanczos2d_aa_backward_kernel_impl) } // namespace at::native diff --git a/aten/src/ATen/native/cpu/UpSampleKernelAVXAntialias.h b/aten/src/ATen/native/cpu/UpSampleKernelAVXAntialias.h index c1bf79dfa44e6..3c351b05dc92f 100644 --- a/aten/src/ATen/native/cpu/UpSampleKernelAVXAntialias.h +++ b/aten/src/ATen/native/cpu/UpSampleKernelAVXAntialias.h @@ -176,7 +176,7 @@ void ImagingResampleHorizontal( // // TODO: we may want to merge that into the fallback code (currently called - // basic_loop_aa_horizontal) + // basic_loop_separable_1d_horizontal) // Although this may not be needed if / when we port all this code to use // Vec.h since this would potentially give us another fall-back implem @@ -252,7 +252,7 @@ void ImagingResampleVertical( // oB[xoffset + i] = b[xoffset + ymin[i]] * w[i, 0] + ... + b[xoffset + ymin[i] + (K-1) * xsize] * w[i, K-1] // TODO: we may want to merge that into the fallback code (currently called - // basic_loop_aa_vertical) + // basic_loop_separable_1d_vertical) // Although this may not be needed if / when we port all this code to use // Vec.h since this would potentially give us another fall-back implem const int16_t* kk = (int16_t*)(vert_indices_weights[3].const_data_ptr()); @@ -289,7 +289,7 @@ void ImagingResampleVertical( // mode for uint8 dtype when C <= 4, with or without antialias. The // implem is based on PIL-SIMD. // Its equivalent implementation (fallback) for when AVX isn't supported or when -// C > 4 is separable_upsample_generic_Nd_kernel_impl() There are a bunch of +// C > 4 is upsample_separable_Nd_kernel_impl() There are a bunch of // future improvement that can be done: look for the TODOs in this file. // For details on how the weights are computed and how the multiplications are // run on int (instead of float weights), see diff --git a/aten/src/ATen/native/cpu/UpSampleKernelNEONAntialias.h b/aten/src/ATen/native/cpu/UpSampleKernelNEONAntialias.h index 6449d27e49f6e..c8d3a1a59a508 100644 --- a/aten/src/ATen/native/cpu/UpSampleKernelNEONAntialias.h +++ b/aten/src/ATen/native/cpu/UpSampleKernelNEONAntialias.h @@ -112,7 +112,13 @@ void NeonResampleHorizontal(const at::Tensor& unpacked_output, } // Block 4: handle 4 pixels that didn't fit in a block of 8 - for (; i + 4 <= ids_size; i += 4) { + // We also use vld3_u8 here, which still loads 8 pixels, but we only use + // the lower half - so the computation is correct. + // On all rows except the last one, reading 8 pixels is safe (the tensors + // are channels-last). But on the last row, we have to be careful not to + // read past the buffer, hence the extra boundary check. + const uint8_t* block_of_4_safe_load_end = input_p + yin * xin_stride - 24; + for (; i + 4 <= ids_size && lineIn_min + num_channels * i <= block_of_4_safe_load_end; i += 4) { uint8x8x3_t rgb = vld3_u8(lineIn_min + num_channels * i); int16x4_t weights4 = vld1_s16(&k[i]); diff --git a/aten/src/ATen/native/cpu/int4mm_kernel.cpp b/aten/src/ATen/native/cpu/int4mm_kernel.cpp index 9a9897b88fba2..72911aea5f5bb 100644 --- a/aten/src/ATen/native/cpu/int4mm_kernel.cpp +++ b/aten/src/ATen/native/cpu/int4mm_kernel.cpp @@ -40,7 +40,7 @@ inline bool is_block_start(int index, int BLOCK_SIZE) { #if (defined(CPU_CAPABILITY_AVX512) || defined(CPU_CAPABILITY_AVX2)) && !defined(_MSC_VER) // convert 16x int4 to int8, handle 64 bits at a time // used in avx2 and avx512 -inline __m128i conver_int4_to_int8(const uint8_t* data) { +inline __m128i convert_int4_to_int8(const uint8_t* data) { __m128i tmp = _mm_loadu_si64((const __m128i*)data); __m128i bytes = _mm_cvtepu8_epi16(tmp); const __m128i lowMask = _mm_set1_epi8(0xF); @@ -169,7 +169,7 @@ inline void tinygemm_kernel( vb[3] = _mm512_fmadd_ps(vb[3], scale[3], zero[3]); } } else { - __m128i b8 = conver_int4_to_int8(B + k * ldb + col * 8); + __m128i b8 = convert_int4_to_int8(B + k * ldb + col * 8); __m512i b32 = _mm512_cvtepu8_epi32(b8); vb[col] = _mm512_permutexvar_ps(b32, lut); vb[col] = _mm512_fmadd_ps(vb[col], scale[col], zero[col]); @@ -312,7 +312,7 @@ inline void tinygemm_kernel( } else { if constexpr (col % 2 == 0) { // de-quantize per 64 bits (16x int4) - __m128i b8 = conver_int4_to_int8(B + k * ldb + col * 4); + __m128i b8 = convert_int4_to_int8(B + k * ldb + col * 4); __m128i b8_val0 = _mm_set1_epi64x(_mm_extract_epi64(b8, 0)); __m128i b8_val1 = _mm_set1_epi64x(_mm_extract_epi64(b8, 1)); if (k + PREFETCH_SIZE_K < K) { diff --git a/aten/src/ATen/native/cuda/AdaptiveMaxPooling2d.cu b/aten/src/ATen/native/cuda/AdaptiveMaxPooling2d.cu index 2030dbb904505..30745cb8f0a0a 100644 --- a/aten/src/ATen/native/cuda/AdaptiveMaxPooling2d.cu +++ b/aten/src/ATen/native/cuda/AdaptiveMaxPooling2d.cu @@ -393,7 +393,7 @@ TORCH_IMPL_FUNC(adaptive_max_pool2d_backward_out_cuda) C10_CUDA_KERNEL_LAUNCH_CHECK(); } else { // run updateGradInput kernel - atomicadaptivemaxgradinput<<< + adaptivemaxgradinput<<< blocks, threads, 0, diff --git a/aten/src/ATen/native/cuda/AveragePool3d.cu b/aten/src/ATen/native/cuda/AveragePool3d.cu index dabcf5b63be99..f290864a9320b 100644 --- a/aten/src/ATen/native/cuda/AveragePool3d.cu +++ b/aten/src/ATen/native/cuda/AveragePool3d.cu @@ -500,11 +500,6 @@ TORCH_IMPL_FUNC(avg_pool3d_backward_out_cuda) ( const int64_t oheight = gradOutput.size(-2); const int64_t owidth = gradOutput.size(-1); - /* XXX shape check behavior from TH */ - const int64_t otime_for_shape_check = pooling_output_shape(itime, kT, padT, dT, 1, ceil_mode); - const int64_t oheight_for_shape_check = pooling_output_shape(iheight, kH, padH, dH, 1, ceil_mode); - const int64_t owidth_for_chape_check = pooling_output_shape(iwidth, kW, padW, dW, 1, ceil_mode); - const bool kernelsOverlap = (dT < kT) || (dH < kH) || (dW < kW); Tensor work_grad_input = gradInput; diff --git a/aten/src/ATen/native/cuda/Blas.cpp b/aten/src/ATen/native/cuda/Blas.cpp index d87395bb480b3..33abc9003e8cb 100644 --- a/aten/src/ATen/native/cuda/Blas.cpp +++ b/aten/src/ATen/native/cuda/Blas.cpp @@ -34,6 +34,7 @@ #else #include #include +#include #include #include #include diff --git a/aten/src/ATen/native/cuda/CUDALoops.cuh b/aten/src/ATen/native/cuda/CUDALoops.cuh index 688ee468939e0..00f0aedf4578c 100644 --- a/aten/src/ATen/native/cuda/CUDALoops.cuh +++ b/aten/src/ATen/native/cuda/CUDALoops.cuh @@ -648,7 +648,6 @@ void gpu_kernel_impl_nocast(TensorIteratorBase& iter, const func_t& f) { TORCH_INTERNAL_ASSERT(iter.can_use_32bit_indexing()); TORCH_INTERNAL_ASSERT(iter.ninputs() == traits::arity); TORCH_INTERNAL_ASSERT(iter.noutputs() == 1); - TORCH_INTERNAL_ASSERT(!needs_dynamic_casting::check(iter)); std::array data; for (int i = 0; i < ntensors; i++) { diff --git a/aten/src/ATen/native/cuda/CUDAScalar.cu b/aten/src/ATen/native/cuda/CUDAScalar.cu index 169a2ab92615f..49af9c8087cdf 100644 --- a/aten/src/ATen/native/cuda/CUDAScalar.cu +++ b/aten/src/ATen/native/cuda/CUDAScalar.cu @@ -10,28 +10,67 @@ #endif #include +#include namespace at::native { +namespace { + +bool is_cuda_caching_allocator_tensor(const Tensor& self) { + auto* cuda_allocator = c10::cuda::CUDACachingAllocator::get(); + if (cuda_allocator == nullptr) { + return false; + } + // SymmMem/NVSHMEM/rocSHMEM tensors are typically backed by custom + // from_blob-style deleters, so this check filters them out and keeps + // the direct dereference path limited to allocator-managed CUDA memory. + return self.storage().data_ptr().get_deleter() == cuda_allocator->raw_deleter(); +} + +template +void _local_scalar_dense_cuda_impl(const Tensor& self, Scalar& r) { +#if defined(USE_ROCM) && (ROCM_VERSION >= 70200) + // If this is a large BAR device, we can just read directly from VRAM + if ( + at::cuda::getCurrentDeviceProperties()->isLargeBar && + is_cuda_caching_allocator_tensor(self)) { + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + hipStreamCaptureStatus captureStatus; + C10_CUDA_CHECK(hipStreamGetCaptureInfo(stream, &captureStatus, nullptr)); + if (C10_LIKELY(captureStatus == hipStreamCaptureStatusNone)) { + at::cuda::stream_synchronize(stream); + r = Scalar(*self.template const_data_ptr()); + } else { + C10_CUDA_CHECK(hipErrorStreamCaptureUnsupported); + } + return; + } +#endif + + // Create pinned memory for the scalar value to avoid implicit + // locking/sync in cuda library due to pageable memory + auto value = at::detail::empty_cpu( + {1}, /* size */ + c10::CppTypeToScalarType(), /* dtype */ + std::nullopt, /* layout */ + std::nullopt, /* device */ + true, /* pin_memory */ + std::nullopt /* memory format */ + ); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + at::cuda::memcpy_and_sync(value.template mutable_data_ptr(), self.template const_data_ptr(), sizeof(scalar_t), cudaMemcpyDeviceToHost, stream); + r = Scalar(*value.template const_data_ptr()); +} + +} // anonymous namespace + Scalar _local_scalar_dense_cuda(const Tensor& self) { Scalar r; TORCH_CHECK(self.numel() > 0, "_local_scalar_dense: Empty tensor not supported"); AT_DISPATCH_V2( self.scalar_type(), "_local_scalar_dense_cuda", AT_WRAP([&] { - // Create pinned memory for the scalar value to avoid implicit - // locking/sync in cuda library due to pageable memory - auto value = at::detail::empty_cpu( - {1}, /* size */ - c10::CppTypeToScalarType(), /* dtype */ - std::nullopt, /* layout */ - std::nullopt, /* device */ - true, /* pin_memory */ - std::nullopt /* memory format */ - ); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - at::cuda::memcpy_and_sync(value.mutable_data_ptr(), self.const_data_ptr(), sizeof(scalar_t), cudaMemcpyDeviceToHost, stream); - r = Scalar(*value.const_data_ptr()); - }), AT_EXPAND(AT_ALL_TYPES_AND_COMPLEX), kComplexHalf, kHalf, kBool, kBFloat16, AT_EXPAND(AT_FLOAT8_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES)); + _local_scalar_dense_cuda_impl(self, r); + }), AT_EXPAND(AT_ALL_TYPES_AND_COMPLEX), kComplexHalf, kHalf, kBool, kBFloat16, AT_EXPAND(AT_FLOAT8_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES)); return r; } diff --git a/aten/src/ATen/native/cuda/CompositeRandomAccessor.h b/aten/src/ATen/native/cuda/CompositeRandomAccessor.h index eb8587d1f9337..d519e333b1e9d 100644 --- a/aten/src/ATen/native/cuda/CompositeRandomAccessor.h +++ b/aten/src/ATen/native/cuda/CompositeRandomAccessor.h @@ -25,7 +25,7 @@ void swap( references_holder rh1, references_holder rh2 ) { - return thrust::swap(rh1.data(), rh2.data()); + thrust::swap(rh1.data(), rh2.data()); } template diff --git a/aten/src/ATen/native/cuda/Copy.cu b/aten/src/ATen/native/cuda/Copy.cu index 3a9abc0588648..68741f90ef85e 100644 --- a/aten/src/ATen/native/cuda/Copy.cu +++ b/aten/src/ATen/native/cuda/Copy.cu @@ -204,7 +204,7 @@ void float8_copy_kernel_cuda(TensorIteratorBase &iter) { break; } } else { - TORCH_CHECK(false, "This supposed ot be called only for Float8 types"); + TORCH_CHECK(false, "This supposed to be called only for Float8 types"); } } diff --git a/aten/src/ATen/native/cuda/Distributions.cpp b/aten/src/ATen/native/cuda/Distributions.cpp index be397f4bc217f..f270ea66d3755 100644 --- a/aten/src/ATen/native/cuda/Distributions.cpp +++ b/aten/src/ATen/native/cuda/Distributions.cpp @@ -28,6 +28,14 @@ Tensor _s_poisson_cuda(const Tensor& lambda, std::optional gen_) { // NOLINTNEXTLINE(performance-unnecessary-value-param) Tensor _s_binomial_cuda(const Tensor& count, const Tensor& prob, std::optional gen_) { + TORCH_CHECK_VALUE( + at::isFloatingType(count.scalar_type()), + "binomial only supports floating-point dtypes for count, got: ", + count.scalar_type()); + TORCH_CHECK_VALUE( + at::isFloatingType(prob.scalar_type()), + "binomial only supports floating-point dtypes for prob, got: ", + prob.scalar_type()); auto gen = get_generator_or_default(gen_, cuda::detail::getDefaultCUDAGenerator()); Tensor ret = at::empty(count.sizes(), count.options()); at::TensorIterator iter = at::TensorIteratorConfig() diff --git a/aten/src/ATen/native/cuda/ForeachFunctors.cuh b/aten/src/ATen/native/cuda/ForeachFunctors.cuh index 37175776097df..1e18e42c34565 100644 --- a/aten/src/ATen/native/cuda/ForeachFunctors.cuh +++ b/aten/src/ATen/native/cuda/ForeachFunctors.cuh @@ -77,6 +77,34 @@ __device__ bool init_args( return all_aligned; } +template < + int depth, + typename param_type, + typename grad_type, + typename exp_avg_type, + typename exp_avg_sq_type> +__device__ bool init_args_mixed_prec( + param_type** param_args, + grad_type** grad_args, + exp_avg_type** exp_avg_args, + exp_avg_sq_type** exp_avg_sq_args, + FusedOptimizerTensorListMetadata& tl, + const int64_t chunk_idx, + const int64_t chunk_size, + const int64_t tensor_loc) { + *param_args = + (param_type*)tl.addresses[0][tensor_loc] + chunk_idx * chunk_size; + *grad_args = (grad_type*)tl.addresses[1][tensor_loc] + chunk_idx * chunk_size; + *exp_avg_args = + (exp_avg_type*)tl.addresses[2][tensor_loc] + chunk_idx * chunk_size; + *exp_avg_sq_args = + (exp_avg_sq_type*)tl.addresses[3][tensor_loc] + chunk_idx * chunk_size; + + bool all_aligned = is_aligned(*param_args) && is_aligned(*grad_args) && + is_aligned(*exp_avg_args) && is_aligned(*exp_avg_sq_args); + return all_aligned; +} + template __device__ void load_args( T r_args[][kILP], @@ -96,6 +124,43 @@ __device__ void load_args( } } +template < + typename T, + typename param_type, + typename grad_type, + typename exp_avg_type, + typename exp_avg_sq_type> +__device__ void load_args( + T r_args[][kILP], + const param_type* param_args, + const grad_type* grad_args, + const exp_avg_type* exp_avg_args, + const exp_avg_sq_type* exp_avg_sq_args, + const int64_t i_start, + const int64_t chunk_size, + const int64_t n) { +#pragma unroll + for (int ii = 0; ii < kILP; ii++) { + const auto i = i_start + threadIdx.x + ii * blockDim.x; + r_args[0][ii] = 0; + if (i < n && i < chunk_size) { + r_args[0][ii] = static_cast(param_args[i]); + } + r_args[1][ii] = 0; + if (i < n && i < chunk_size) { + r_args[1][ii] = static_cast(grad_args[i]); + } + r_args[2][ii] = 0; + if (i < n && i < chunk_size) { + r_args[2][ii] = static_cast(exp_avg_args[i]); + } + r_args[3][ii] = 0; + if (i < n && i < chunk_size) { + r_args[3][ii] = static_cast(exp_avg_sq_args[i]); + } + } +} + template __device__ void store_args( T* dst, @@ -111,6 +176,21 @@ __device__ void store_args( } } +template +__device__ void store_args( + dT* dst, + sT* src, + const int64_t i_start, + const int64_t chunk_size, + const int64_t n) { +#pragma unroll + for (int ii = 0; ii < kILP; ii++) { + const int64_t i = i_start + threadIdx.x + ii * blockDim.x; + if (i < n && i < chunk_size) + dst[i] = static_cast
(src[ii]); + } +} + template __device__ __forceinline__ void binary_op_scalar( T r_args[][kILP], diff --git a/aten/src/ATen/native/cuda/ForeachReduceOp.cu b/aten/src/ATen/native/cuda/ForeachReduceOp.cu index 78f0434d9f7a8..e7a073fe75a55 100644 --- a/aten/src/ATen/native/cuda/ForeachReduceOp.cu +++ b/aten/src/ATen/native/cuda/ForeachReduceOp.cu @@ -28,22 +28,25 @@ namespace at::native { -// _foreach_norm supports only L1, L2, and inf norm -enum class NormType { L1, L2, LInf }; +// _foreach_norm supports L0, L1, L2, and inf norm +enum class NormType { L0, L1, L2, LInf }; // NOTE: This is a simple variant of TensorListMetadata in MultiTensorApply.cuh // as we only need to track addresses for the lpnorm_cleanup function below. // Why is this struct necessary? For the same reason the TensorListMetadata // struct is necessary--which is to ferry static metadata to the CUDA kernel -// while complying with the 4kb size constraint. Since we only need to track -// addresses, we introduce this struct to be able to fit more Tensor pointers at -// a time, currently 400 empirically, compared to the much smaller values in -// depth_to_max_tensors. This way, we can launch fewer kernels for better -// performance. +// while complying with the kernel arg size constraint. Since we only need to +// track addresses, we introduce this struct to be able to fit more Tensor +// pointers at a time compared to depth_to_max_tensors. This way, we can +// launch fewer kernels for better performance. // // IF YOU USE THIS STRUCT, PLEASE ADD A ONE-OFF TEST IN test_foreach.py AS THIS // IS CURRENTLY ONLY TESTED FOR _foreach_norm. -const size_t MAX_TENSORS_PER_KERNEL = 400; +#if defined(CUDART_VERSION) && CUDART_VERSION >= 13000 && !defined(USE_ROCM) +static constexpr size_t MAX_TENSORS_PER_KERNEL = 3200; +#else +static constexpr size_t MAX_TENSORS_PER_KERNEL = 400; +#endif struct TensorListAddresses { const void* addresses[MAX_TENSORS_PER_KERNEL]; }; @@ -281,7 +284,10 @@ struct LpNormFunctor { #pragma unroll for (int ii = 0; ii < kILP; ii++) { const auto next = static_cast(r_x[ii]); - if constexpr (norm_type == NormType::LInf) { + if constexpr (norm_type == NormType::L0) { + vals[ii] += + next != out_opmath_t(0) ? out_opmath_t(1) : out_opmath_t(0); + } else if constexpr (norm_type == NormType::LInf) { vals[ii] = max_propagate_nan(vals[ii], ::abs(next)); } else { vals[ii] += norm_type == NormType::L1 ? ::abs(next) : next * next; @@ -296,7 +302,10 @@ struct LpNormFunctor { int i = i_start + threadIdx.x + ii * blockDim.x; if (i < n && i < chunk_size) { const auto next = static_cast(x[i]); - if constexpr (norm_type == NormType::LInf) { + if constexpr (norm_type == NormType::L0) { + vals[ii] += + next != out_opmath_t(0) ? out_opmath_t(1) : out_opmath_t(0); + } else if constexpr (norm_type == NormType::LInf) { vals[ii] = max_propagate_nan(vals[ii], ::abs(next)); } else { vals[ii] += norm_type == NormType::L1 ? ::abs(next) : next * next; @@ -314,7 +323,8 @@ struct LpNormFunctor { val += vals[i]; } } - auto final_val = norm_type == NormType::L1 || norm_type == NormType::L2 + auto final_val = norm_type == NormType::L0 || norm_type == NormType::L1 || + norm_type == NormType::L2 ? at::native::cuda_utils::BlockReduceSum(val, s_vals) : at::native::cuda_utils::BlockReduceMax(val, s_vals); @@ -348,8 +358,8 @@ __global__ void lpnorm_cleanup( val += output_this_tensor[i]; } } - out_opmath_t final_val = - norm_type == NormType::L1 || norm_type == NormType::L2 + out_opmath_t final_val = norm_type == NormType::L0 || + norm_type == NormType::L1 || norm_type == NormType::L2 ? at::native::cuda_utils::BlockReduceSum(val, vals) : at::native::cuda_utils::BlockReduceMax(val, vals); if (threadIdx.x == 0) { @@ -476,7 +486,13 @@ std::vector foreach_tensor_norm_cuda_internal( AT_DISPATCH_OUT_DTYPES( output_dtype, ForeachNormDispatchName::value, [&]() { using out_opmath_t = typename at::opmath_type; - if (p == static_cast(1)) { + if (p == static_cast(0)) { + multi_tensor_apply<1>( + tensor_lists, + LpNormFunctor(), + output_per_tensor.template mutable_data_ptr(), + max_chunks_per_tensor); + } else if (p == static_cast(1)) { multi_tensor_apply<1>( tensor_lists, LpNormFunctor(), @@ -519,7 +535,16 @@ std::vector foreach_tensor_norm_cuda_internal( .template mutable_data_ptr(); } - if (p == static_cast(1)) { + if (p == static_cast(0)) { + lpnorm_cleanup + <<>>( + output_per_tensor + .template const_data_ptr() + + i * MAX_TENSORS_PER_KERNEL * + max_chunks_per_tensor, + addr_struct, + max_chunks_per_tensor); + } else if (p == static_cast(1)) { lpnorm_cleanup <<>>( output_per_tensor @@ -613,7 +638,8 @@ std::vector foreach_tensor_norm_cuda( }); } if (!can_use_fast_route(tensors) || has_int_or_complex || - !(p == static_cast(1) || p == static_cast(2) || + !(p == static_cast(0) || p == static_cast(1) || + p == static_cast(2) || p == std::numeric_limits::infinity())) { return foreach_tensor_norm_slow(tensors, ord, dtype); } diff --git a/aten/src/ATen/native/cuda/ForeachUnaryOp.cu b/aten/src/ATen/native/cuda/ForeachUnaryOp.cu index bb070f9d97616..9a36a879615a0 100644 --- a/aten/src/ATen/native/cuda/ForeachUnaryOp.cu +++ b/aten/src/ATen/native/cuda/ForeachUnaryOp.cu @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -40,7 +41,9 @@ #include #include +#include #include +#include #endif namespace at::native { @@ -431,4 +434,48 @@ void foreach_tensor_zero_cuda_(TensorList tensors) { }); } +std::vector foreach_tensor_clone_cuda( + TensorList self, + std::optional memory_format) { + check_foreach_api_restrictions(self); + if (!_check_tensors_share_device_and_dtype({self})) { + return at::native::foreach_tensor_clone_slow(self, memory_format); + } + + std::vector ret{}; + ret.reserve(self.size()); + + auto realized_memory_format = memory_format.value_or(MemoryFormat::Preserve); + for (const auto& s : self) { + // This logic modified from at::native::clone. + if (realized_memory_format == MemoryFormat::Preserve) { + if (s.is_non_overlapping_and_dense()) { + // Copy all strides, this is marginally faster than calling empty_like + auto options = s.options(); + ret.emplace_back(at::native::empty_strided_cuda( + s.sizes(), + s.strides(), + c10::optTypeMetaToScalarType(options.dtype_opt()), + options.layout_opt(), + options.device_opt(), + options.pinned_memory_opt())); + } else { + ret.emplace_back(at::native::empty_like(s)); + } + } else { + auto options = s.options(); + ret.emplace_back(at::native::empty_like( + s, + c10::optTypeMetaToScalarType(options.dtype_opt()), + options.layout_opt(), + options.device_opt(), + options.pinned_memory_opt(), + realized_memory_format)); + } + } + + at::native::foreach_tensor_copy_list_kernel_cuda_(ret, self); + return ret; +} + } // namespace at::native diff --git a/aten/src/ATen/native/cuda/FusedAdamKernel.cu b/aten/src/ATen/native/cuda/FusedAdamKernel.cu index 0858f24e17c6e..498dc634fab2b 100644 --- a/aten/src/ATen/native/cuda/FusedAdamKernel.cu +++ b/aten/src/ATen/native/cuda/FusedAdamKernel.cu @@ -29,10 +29,15 @@ void _fused_adam_kernel_cuda_( const bool maximize, const std::optional& grad_scale, const std::optional& found_inf) { + const bool is_mixed_precision = + params[0].scalar_type() != exp_avgs[0].scalar_type(); if (amsgrad) { TORCH_CHECK( at::native::check_fast_path_restrictions( - {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}), + {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}, + /*scalarList=*/{}, + /*does_op_promote_integer_inputs_to_float=*/false, + /*skip_cross_list_dtype_check=*/is_mixed_precision), "params, grads, exp_avgs, exp_avg_sqs, and max_exp_avg_sqs must have same dtype, device, and layout"); _fused_adam_amsgrad_cuda_impl_( params, @@ -52,7 +57,10 @@ void _fused_adam_kernel_cuda_( } else { TORCH_CHECK( at::native::check_fast_path_restrictions( - {params, grads, exp_avgs, exp_avg_sqs}), + {params, grads, exp_avgs, exp_avg_sqs}, + /*scalarList=*/{}, + /*does_op_promote_integer_inputs_to_float=*/false, + /*skip_cross_list_dtype_check=*/is_mixed_precision), "params, grads, exp_avgs, and exp_avg_sqs must have same dtype, device, and layout"); _fused_adam_cuda_impl_( params, @@ -125,10 +133,15 @@ void _fused_adam_kernel_cuda_( lr.device() == param_device, "lr must be on the same GPU device as the params"); + const bool is_mixed_precision = + params[0].scalar_type() != exp_avgs[0].scalar_type(); if (amsgrad) { TORCH_CHECK( at::native::check_fast_path_restrictions( - {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}), + {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}, + /*scalarList=*/{}, + /*does_op_promote_integer_inputs_to_float=*/false, + /*skip_cross_list_dtype_check=*/is_mixed_precision), "params, grads, exp_avgs, exp_avg_sqs, and max_exp_avg_sqs must have same dtype, device, and layout"); _fused_adam_amsgrad_cuda_impl_( params, @@ -148,7 +161,10 @@ void _fused_adam_kernel_cuda_( } else { TORCH_CHECK( at::native::check_fast_path_restrictions( - {params, grads, exp_avgs, exp_avg_sqs}), + {params, grads, exp_avgs, exp_avg_sqs}, + /*scalarList=*/{}, + /*does_op_promote_integer_inputs_to_float=*/false, + /*skip_cross_list_dtype_check=*/is_mixed_precision), "params, grads, exp_avgs, and exp_avg_sqs must have same dtype, device, and layout"); _fused_adam_cuda_impl_( params, diff --git a/aten/src/ATen/native/cuda/FusedAdamWKernel.cu b/aten/src/ATen/native/cuda/FusedAdamWKernel.cu index 4b758aa574c66..13e11df2a24e9 100644 --- a/aten/src/ATen/native/cuda/FusedAdamWKernel.cu +++ b/aten/src/ATen/native/cuda/FusedAdamWKernel.cu @@ -29,10 +29,15 @@ void _fused_adamw_kernel_cuda_( const bool maximize, const std::optional& grad_scale, const std::optional& found_inf) { + const bool is_mixed_precision = + params[0].scalar_type() != exp_avgs[0].scalar_type(); if (amsgrad) { TORCH_CHECK( at::native::check_fast_path_restrictions( - {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}), + {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}, + /*scalarList=*/{}, + /*does_op_promote_integer_inputs_to_float=*/false, + /*skip_cross_list_dtype_check=*/is_mixed_precision), "params, grads, exp_avgs, exp_avg_sqs, and max_exp_avg_sqs must have same dtype, device, and layout"); _fused_adamw_amsgrad_cuda_impl_( params, @@ -52,7 +57,10 @@ void _fused_adamw_kernel_cuda_( } else { TORCH_CHECK( at::native::check_fast_path_restrictions( - {params, grads, exp_avgs, exp_avg_sqs}), + {params, grads, exp_avgs, exp_avg_sqs}, + /*scalarList=*/{}, + /*does_op_promote_integer_inputs_to_float=*/false, + /*skip_cross_list_dtype_check=*/is_mixed_precision), "params, grads, exp_avgs, and exp_avg_sqs must have same dtype, device, and layout"); _fused_adamw_cuda_impl_( params, @@ -125,10 +133,15 @@ void _fused_adamw_kernel_cuda_( lr.device() == param_device, "lr must be on the same GPU device as the params"); + const bool is_mixed_precision = + params[0].scalar_type() != exp_avgs[0].scalar_type(); if (amsgrad) { TORCH_CHECK( at::native::check_fast_path_restrictions( - {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}), + {params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs}, + /*scalarList=*/{}, + /*does_op_promote_integer_inputs_to_float=*/false, + /*skip_cross_list_dtype_check=*/is_mixed_precision), "params, grads, exp_avgs, exp_avg_sqs, and max_exp_avg_sqs must have same dtype, device, and layout"); _fused_adamw_amsgrad_cuda_impl_( params, @@ -148,7 +161,10 @@ void _fused_adamw_kernel_cuda_( } else { TORCH_CHECK( at::native::check_fast_path_restrictions( - {params, grads, exp_avgs, exp_avg_sqs}), + {params, grads, exp_avgs, exp_avg_sqs}, + /*scalarList=*/{}, + /*does_op_promote_integer_inputs_to_float=*/false, + /*skip_cross_list_dtype_check=*/is_mixed_precision), "params, grads, exp_avgs, and exp_avg_sqs must have same dtype, device, and layout"); _fused_adamw_cuda_impl_( params, diff --git a/aten/src/ATen/native/cuda/GroupedBlas.cpp b/aten/src/ATen/native/cuda/GroupedBlas.cpp index 8c9125879dc06..5b393401a29b4 100644 --- a/aten/src/ATen/native/cuda/GroupedBlas.cpp +++ b/aten/src/ATen/native/cuda/GroupedBlas.cpp @@ -37,6 +37,7 @@ #else #include #include +#include #include #include #include @@ -152,7 +153,7 @@ _f8_f8_bf16_rowwise_grouped_mm_cuda( const bool use_fast_accum, Tensor& out) { TORCH_CHECK_VALUE(mat_a.dtype() == at::kFloat8_e4m3fn, "Expected mat_a to be Float8_e4m3 matrix got ", mat_a.scalar_type()); - TORCH_CHECK_VALUE(mat_b.dtype() == at::kFloat8_e4m3fn, "Expected mat_a to be Float8_e4m3 matrix got ", mat_b.scalar_type()); + TORCH_CHECK_VALUE(mat_b.dtype() == at::kFloat8_e4m3fn, "Expected mat_b to be Float8_e4m3 matrix got ", mat_b.scalar_type()); at::cuda::detail::f8f8bf16_grouped_mm( mat_a, diff --git a/aten/src/ATen/native/cuda/IndexKernelUtils.cu b/aten/src/ATen/native/cuda/IndexKernelUtils.cu index e44450c542e83..d9bf529f38db1 100644 --- a/aten/src/ATen/native/cuda/IndexKernelUtils.cu +++ b/aten/src/ATen/native/cuda/IndexKernelUtils.cu @@ -30,7 +30,7 @@ void vectorized_gather_kernel_launch(char * out, char * inp, index_t * idx, int constexpr int64_t max_num_threads=256; auto num_threads = at::round_up( at::ceil_div(slice_size_in_bytes, Alignment), - static_cast(C10_WARP_SIZE)); + static_cast(at::cuda::warp_size())); uint32_t grid_y = at::cuda::getCurrentDeviceProperties()->maxGridSize[1]; grid_y = std::min(static_cast(at::ceil_div(slice_size_in_bytes, max_num_threads * Alignment)), grid_y); dim3 grid = {static_cast(num_ind), grid_y, 1}; diff --git a/aten/src/ATen/native/cuda/Loops.cuh b/aten/src/ATen/native/cuda/Loops.cuh index 962f71431e727..063db29d38279 100644 --- a/aten/src/ATen/native/cuda/Loops.cuh +++ b/aten/src/ATen/native/cuda/Loops.cuh @@ -81,7 +81,7 @@ __device__ inline void elementwise_kernel_helper(func_t f, policy_t policy) { namespace at:: native { template -void gpu_kernel_nocast(TensorIteratorBase& iter, const func_t& f) { +void gpu_kernel_nocast(TensorIteratorBase& iter, const func_t& f, bool check_cast = true) { for (int arg = 0; arg < iter.ntensors(); arg++) { TORCH_INTERNAL_ASSERT( @@ -95,14 +95,22 @@ void gpu_kernel_nocast(TensorIteratorBase& iter, const func_t& f) { if (!iter.can_use_32bit_indexing()) { for (auto& sub_iter : iter.with_32bit_indexing()) { - gpu_kernel_nocast(sub_iter, f); + gpu_kernel_nocast(sub_iter, f, check_cast); } return; } + if (check_cast) { + TORCH_INTERNAL_ASSERT(!needs_dynamic_casting::check(iter)); + } gpu_kernel_impl_nocast(iter, f); } +template +void gpu_kernel_opaque(TensorIteratorBase& iter, const func_t& f) { + gpu_kernel_nocast(iter, f, false); +} + template void gpu_kernel(TensorIteratorBase& iter, const func_t& f) { diff --git a/aten/src/ATen/native/cuda/MultiTensorApply.cuh b/aten/src/ATen/native/cuda/MultiTensorApply.cuh index 2fe431f778b1a..d11cfc19c5856 100644 --- a/aten/src/ATen/native/cuda/MultiTensorApply.cuh +++ b/aten/src/ATen/native/cuda/MultiTensorApply.cuh @@ -15,13 +15,35 @@ static constexpr int64_t kChunkSize = 65536; static constexpr int64_t kBlockSize = 512; // TODO(crcrpar): Add `n>5` for `low prec params & their higher prec copy` -// TensorListMetadata has to be < 4KB - the limit for kernel launch argument +// TensorListMetadata has to fit within the CUDA kernel launch argument limit. +// While CUDA 12.1, driver version R530+ and Volta+ would work with 32KB, we +// decide to be safe and only swap for CUDA 13+ during compile time. This saves +// binary size and will guarantees 32KB kernel arg space; older versions are +// still limited to 4KB. We adopt naive values for 32KB from +// https://github.com/pytorch/pytorch/pull/134373. +// TODO: The values for 32KB can very much be optimized further. +#if defined(CUDART_VERSION) && CUDART_VERSION >= 13000 && !defined(USE_ROCM) + +static constexpr int depth_to_max_tensors[5] = {770, 448, 336, 252, 210}; +static constexpr int depth_to_max_blocks[5] = {2240, 2240, 2240, 2240, 2240}; +static constexpr int depth_to_max_tensors_scalarlist[5] = + {672, 448, 336, 252, 210}; +static constexpr int depth_to_max_tensors_scalarlist_of_complex_double[2] = { + 504, + 420}; +using block_index_t = uint16_t; + +#else + static constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30}; static constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320}; static constexpr int depth_to_max_tensors_scalarlist[5] = {96, 64, 48, 36, 30}; static constexpr int depth_to_max_tensors_scalarlist_of_complex_double[2] = { 72, 60}; +using block_index_t = unsigned char; + +#endif template __device__ __forceinline__ bool is_aligned(T* p) { @@ -42,7 +64,7 @@ template struct TensorListMetadata { const void* addresses[n][depth_to_max_tensors[n - 1]]; int64_t numel_for_tensor[depth_to_max_tensors[n - 1]]; - unsigned char block_to_tensor[depth_to_max_blocks[n - 1]]; + block_index_t block_to_tensor[depth_to_max_blocks[n - 1]]; int block_to_chunk[depth_to_max_blocks[n - 1]]; int start_tensor_this_launch; }; @@ -52,12 +74,12 @@ struct TensorListScalarListMetadata { const void* addresses[n][depth_to_max_tensors_scalarlist[n - 1]]; int64_t numel_for_tensor[depth_to_max_tensors_scalarlist[n - 1]]; scalar_vals_t scalar_vals[depth_to_max_tensors_scalarlist[n - 1]]; - unsigned char block_to_tensor[depth_to_max_blocks[n - 1]]; + block_index_t block_to_tensor[depth_to_max_blocks[n - 1]]; int block_to_chunk[depth_to_max_blocks[n - 1]]; }; -// note(mkozuki): `n` of 1&2 violate the limit of cuda kernel argument size of -// 4kb with `c10::complex` +// note(mkozuki): `n` of 1&2 violate the limit of cuda kernel argument size +// with `c10::complex` template <> struct TensorListScalarListMetadata, 1> { const void* addresses[1] @@ -66,7 +88,7 @@ struct TensorListScalarListMetadata, 1> { numel_for_tensor[depth_to_max_tensors_scalarlist_of_complex_double[0]]; c10::complex scalar_vals[depth_to_max_tensors_scalarlist_of_complex_double[0]]; - unsigned char block_to_tensor[depth_to_max_blocks[1 - 1]]; + block_index_t block_to_tensor[depth_to_max_blocks[1 - 1]]; int block_to_chunk[depth_to_max_blocks[1 - 1]]; }; @@ -78,19 +100,22 @@ struct TensorListScalarListMetadata, 2> { numel_for_tensor[depth_to_max_tensors_scalarlist_of_complex_double[1]]; c10::complex scalar_vals[depth_to_max_tensors_scalarlist_of_complex_double[1]]; - unsigned char block_to_tensor[depth_to_max_blocks[2 - 1]]; + block_index_t block_to_tensor[depth_to_max_blocks[2 - 1]]; int block_to_chunk[depth_to_max_blocks[2 - 1]]; }; // NOTE(crcrpar): This is a conservative resolution to handle `state_steps` // whose each element is `at::Tensor` of 1 element representing the number of // `step`s called so far. +// We're aware this struct overflows the kernel arg limit at n=1 (4244 bytes), +// but our current fused optimizers only instantiate at n>=4 so it's not a +// concern (yet). template struct FusedOptimizerTensorListMetadata { const void* addresses[n][depth_to_max_tensors[n - 1]]; int64_t numel_for_tensor[depth_to_max_tensors[n - 1]]; - const void* state_steps_addresses[depth_to_max_tensors_scalarlist[n - 1]]; - unsigned char block_to_tensor[depth_to_max_blocks[n - 1]]; + const void* state_steps_addresses[depth_to_max_tensors[n - 1]]; + block_index_t block_to_tensor[depth_to_max_blocks[n - 1]]; int block_to_chunk[depth_to_max_blocks[n - 1]]; int start_tensor_this_launch; }; diff --git a/aten/src/ATen/native/cuda/Normalization.cuh b/aten/src/ATen/native/cuda/Normalization.cuh index 8e31f8fa9a694..d9ceeea53ceec 100644 --- a/aten/src/ATen/native/cuda/Normalization.cuh +++ b/aten/src/ATen/native/cuda/Normalization.cuh @@ -134,7 +134,7 @@ __device__ scalar_t reduce(Op op, PTA tensor, int plane) { } #endif } - __shared__ scalar_t shared[C10_WARP_SIZE]; + __shared__ scalar_t shared[C10_WARP_SIZE_UPPER_BOUND]; SumReduceOp reduce_op; sum = cuda_utils::BlockReduce, cuda_utils::Block2D>(sum, reduce_op, 0, shared); if (threadIdx.x == 0 && threadIdx.y == 0) { @@ -288,7 +288,7 @@ __global__ void batch_norm_collect_statistics_kernel( GenericPackedTensorAccessor save_mean, GenericPackedTensorAccessor save_transformed_var) { - __shared__ int shared_n[2 * 2 * C10_WARP_SIZE + C10_WARP_SIZE]; + __shared__ int shared_n[2 * 2 * C10_WARP_SIZE_UPPER_BOUND + C10_WARP_SIZE_UPPER_BOUND]; int plane = blockIdx.x; int N = input.size(0) * input.size(2); @@ -750,7 +750,7 @@ void batch_norm_elemt_cuda_template(const Tensor& output_, const Tensor& input_, // The input_transform kernel is pointwise, but we need to balance reading parameters (save_var/mean, // weight/bias) - which we only do once and have a for loop afterwards - with having many threads and blocks - // and good occupancy. Quiet likely, we could go with even more blocks than 1024. + // and good occupancy. Quite likely, we could go with even more blocks than 1024. // The various planes are independent, so we use blocks for them. int tf = std::max(getNumThreads(input.size(2)/4), std::min(getNumThreads(input.size(2)), 64)); diff --git a/aten/src/ATen/native/cuda/PersistentSoftmax.cuh b/aten/src/ATen/native/cuda/PersistentSoftmax.cuh index f0871fa0ead6f..fee3d0a5aba1d 100644 --- a/aten/src/ATen/native/cuda/PersistentSoftmax.cuh +++ b/aten/src/ATen/native/cuda/PersistentSoftmax.cuh @@ -64,12 +64,12 @@ __device__ __forceinline__ void warp_reduce(acc_t* sum) { // input_t=half, acc_t=float, output_t=float => read half tensor, float accumulators, write float tensor. // input_t_float, acc_t=float, output_t=half => read float tensor, float accumulators, write half tensor. -template +template __global__ void softmax_warp_forward(output_t *dst, const input_t *src, int batch_size, int stride, int element_count, const bool *mask = nullptr, const int head_chunk_size = -1, bool is_transformer_mask = false) { // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method warp_softmax_forward_kernel. constexpr int next_power_of_two = 1 << log2_elements; - constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE; + constexpr int WARP_SIZE = (next_power_of_two < WARP_SIZE_PARAM) ? next_power_of_two : WARP_SIZE_PARAM; constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE; constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1; @@ -211,12 +211,12 @@ __global__ void softmax_warp_forward(output_t *dst, const input_t *src, int batc } } -template +template __global__ void softmax_warp_backward(output_t *gradInput, const input_t *grad, const input_t *output, int batch_size, int stride, int element_count, const bool *mask = nullptr) { // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method warp_softmax_backward_kernel. constexpr int next_power_of_two = 1 << log2_elements; - constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE; + constexpr int WARP_SIZE = (next_power_of_two < WARP_SIZE_PARAM) ? next_power_of_two : WARP_SIZE_PARAM; constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE; constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1; @@ -325,12 +325,29 @@ void dispatch_softmax_forward(output_t *dst, const input_t *src, int softmax_ele dim3 threads(warp_size, warps_per_block, 1); // Launch code would be more elegant if C++ supported FOR CONSTEXPR switch (log2_elements) { +#ifdef USE_ROCM + // To support ROCm amdgcnspirv target, we must compile both a 32 and 64 warpSize version of each kernel #define LAUNCH_SOFTMAX_WARP_FORWARD(L2E) case L2E: \ - softmax_warp_forward \ + if (warp_size == 64) { \ + softmax_warp_forward \ + <<>>(dst, \ + src, batch_count, softmax_elements_stride, softmax_elements, mask, chunk_size, is_transformer_mask); \ + } \ + else { \ + softmax_warp_forward \ + <<>>(dst, \ + src, batch_count, softmax_elements_stride, softmax_elements, mask, chunk_size, is_transformer_mask); \ + } \ + C10_CUDA_KERNEL_LAUNCH_CHECK(); \ + break; +#else + #define LAUNCH_SOFTMAX_WARP_FORWARD(L2E) case L2E: \ + softmax_warp_forward \ <<>>(dst, \ src, batch_count, softmax_elements_stride, softmax_elements, mask, chunk_size, is_transformer_mask); \ C10_CUDA_KERNEL_LAUNCH_CHECK(); \ break; +#endif LAUNCH_SOFTMAX_WARP_FORWARD(0); // 1 LAUNCH_SOFTMAX_WARP_FORWARD(1); // 2 @@ -376,13 +393,32 @@ void dispatch_softmax_backward(output_t *grad_input, const input_t *grad, const dim3 threads(warp_size, warps_per_block, 1); // Launch code would be more elegant if C++ supported FOR CONSTEXPR switch (log2_elements) { +#ifdef USE_ROCM + // To support ROCm amdgcnspirv target, we must compile both a 32 and 64 warpSize version of each kernel + #define LAUNCH_SOFTMAX_WARP_BACKWARD(L2E) case L2E: \ + if (warp_size == 64) { \ + softmax_warp_backward \ + <<>> \ + (grad_input, grad, output, batch_count, softmax_elements_stride, \ + softmax_elements, mask); \ + } \ + else { \ + softmax_warp_backward \ + <<>> \ + (grad_input, grad, output, batch_count, softmax_elements_stride, \ + softmax_elements, mask); \ + } \ + C10_CUDA_KERNEL_LAUNCH_CHECK(); \ + break; +#else #define LAUNCH_SOFTMAX_WARP_BACKWARD(L2E) case L2E: \ - softmax_warp_backward \ + softmax_warp_backward \ <<>> \ (grad_input, grad, output, batch_count, softmax_elements_stride, \ softmax_elements, mask); \ C10_CUDA_KERNEL_LAUNCH_CHECK(); \ break; +#endif LAUNCH_SOFTMAX_WARP_BACKWARD(0); // 1 LAUNCH_SOFTMAX_WARP_BACKWARD(1); // 2 diff --git a/aten/src/ATen/native/cuda/PhiloxDistribution.cu b/aten/src/ATen/native/cuda/PhiloxDistribution.cu new file mode 100644 index 0000000000000..193ea9618e748 --- /dev/null +++ b/aten/src/ATen/native/cuda/PhiloxDistribution.cu @@ -0,0 +1,322 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +#include +#include +#endif + +namespace at::native { + +namespace { + +using at::cuda::philox_4x32; + +// Elements produced per Philox 4x32 call: 4 for float/half/bfloat16, 2 for double. +// Note that we use a full float for each generated half/bfloat16 for better numerics. +template +constexpr int elems_per_call = std::is_same_v ? 2 : 4; + +// Box-Muller: convert 4 uniform uint32 values into 4 standard normal floats. +__device__ __forceinline__ float4 box_muller_float(uint4 r) { + constexpr float M = 2.3283064365386963e-10f; // 1/2^32 + constexpr float TWO_PI = 6.2831853071795864f; + // Map to (0, 1] to avoid log(0). + float u1 = fmaf(r.x, M, M * 0.5f); + float u2 = fmaf(r.y, M, M * 0.5f); + float u3 = fmaf(r.z, M, M * 0.5f); + float u4 = fmaf(r.w, M, M * 0.5f); + + float radius1 = sqrtf(-2.0f * __logf(u1)); + float radius2 = sqrtf(-2.0f * __logf(u3)); + float s1, c1, s2, c2; + __sincosf(TWO_PI * u2, &s1, &c1); + __sincosf(TWO_PI * u4, &s2, &c2); + return {radius1 * c1, radius1 * s1, radius2 * c2, radius2 * s2}; +} + +// Box-Muller: convert 4 uint32 values (packed into 2 uint64) into 2 standard +// normal doubles. +__device__ __forceinline__ double2 box_muller_double(uint4 r) { + constexpr double M = 2.3283064365386963e-10; // 1/2^32 + constexpr double TWO_PI = 6.2831853071795864; + // Pack pairs of uint32 for ~64 bits of uniform randomness. + double u1 = fma(static_cast(r.x), M, + static_cast(r.y) * M * M + M * M * 0.5); + double u2 = fma(static_cast(r.z), M, + static_cast(r.w) * M * M + M * M * 0.5); + + double radius = ::sqrt(-2.0 * ::log(u1)); + double s, c; + ::sincos(TWO_PI * u2, &s, &c); + return {radius * c, radius * s}; +} + +// Single-key kernel: one thread per chunk of elements, where each chunk +// comes from a single Philox 4x32 call. Uses vectorized stores for full +// chunks and scalar writes for the tail. +template +__global__ void philox_single_key_kernel( + scalar_t* __restrict__ output, + const uint64_t* __restrict__ key, + int64_t num_elems, + sample_t sample_func, + param_t param_func) { + + // Use vectorized load to get (seed, offset) + auto key_vec = memory::ld_vec<16>(key); + auto* key_vals = reinterpret_cast(&key_vec); + uint64_t seed = key_vals[0]; + uint64_t offset = key_vals[1]; + + // Use vectorized stores for full chunks since they're aligned. + constexpr int epc = elems_per_call; + int64_t num_full_chunks = num_elems / epc; + int64_t chunk = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (chunk < num_full_chunks) { + auto sample = sample_func(seed, offset + static_cast(chunk)); + constexpr int vec_bytes = epc * sizeof(scalar_t); + memory::Vec v; + auto* vals = reinterpret_cast(&v); + #pragma unroll + for (int j = 0; j < epc; j++) { + vals[j] = param_func((&sample.x)[j]); + } + memory::st_vec(output + chunk * epc, v); + } + + // Scalar tail for remaining elements. + if (chunk == num_full_chunks) { + int64_t tail_start = num_full_chunks * epc; + auto sample = sample_func(seed, offset + static_cast(num_full_chunks)); + for (int j = 0; j < num_elems - tail_start; j++) { + output[tail_start + j] = param_func((&sample.x)[j]); + } + } +} + +// Multi-key kernel: one thread per (key_idx, chunk) pair, where each chunk +// comes from a single Philox 4x32 call. Uses vectorized stores for full +// chunks and scalar writes for the tail. +template +__global__ void philox_multi_key_kernel( + scalar_t* __restrict__ output, + const uint64_t* __restrict__ keys, + int64_t num_keys, + int64_t elems_per_key, + sample_t sample_func, + param_t param_func, + OffsetCalculator<1> key_offset_calc) { + constexpr int epc = elems_per_call; + int64_t chunks_per_key = (elems_per_key + epc - 1) / epc; + int64_t total_threads = num_keys * chunks_per_key; + int64_t tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (tid >= total_threads) return; + + // Determine correct (seed, offset) to use and sample. + int64_t key_idx = tid / chunks_per_key; + int64_t chunk = tid % chunks_per_key; + auto elem_offset = key_offset_calc.get(key_idx)[0]; + uint64_t seed = keys[elem_offset]; + uint64_t offset = keys[elem_offset + 1]; + auto sample = sample_func(seed, offset + static_cast(chunk)); + + // Vectorized writes require aligned base addresses. This is guaranteed + // when elems_per_key is a multiple of epc, since + // base = key_idx * elems_per_key + chunk * epc. + int64_t full_chunks_per_key = elems_per_key / epc; + bool aligned = elems_per_key % epc == 0; + int64_t base = key_idx * elems_per_key + chunk * epc; + if (aligned && chunk < full_chunks_per_key) { + constexpr int vec_bytes = epc * sizeof(scalar_t); + memory::Vec v; + auto* vals = reinterpret_cast(&v); + #pragma unroll + for (int j = 0; j < epc; j++) { + vals[j] = param_func((&sample.x)[j]); + } + memory::st_vec(output + base, v); + } else { + for (int j = 0; j < epc && chunk * epc + j < elems_per_key; j++) { + output[base + j] = param_func((&sample.x)[j]); + } + } +} + +// Dispatches to single-key or multi-key kernels as needed. +template +void philox_distribution_kernel( + const char* op_name, + Tensor& self, const Tensor& key, + const sample_t& sample_func, const param_t& param_func) { + TORCH_CHECK(self.is_floating_point(), + op_name, ": self must be a floating point tensor, got ", + self.scalar_type()); + TORCH_CHECK(key.scalar_type() == kUInt64, + op_name, ": key must have dtype uint64, got ", + key.scalar_type()); + TORCH_CHECK(self.device() == key.device(), + op_name, ": self and key must be on the same device, got ", + self.device(), " and ", key.device()); + TORCH_CHECK(key.dim() >= 1 && key.size(-1) == 2, + op_name, ": key must have shape (2,) or (*batch, 2), got shape ", + key.sizes()); + if (key.dim() > 1) { + TORCH_CHECK(key.dim() == self.dim() + 1, + op_name, ": batched key must have ndim == output ndim + 1, " + "got key shape ", key.sizes(), " with output shape ", self.sizes()); + auto key_batch = key.sizes().slice(0, self.dim()); + TORCH_CHECK(is_expandable_to(key_batch, self.sizes()), + op_name, ": key batch shape ", key_batch, + " is not broadcastable with output shape ", self.sizes()); + } + + if (self.numel() == 0) { + return; + } + + // Ensure contiguous, aligned output for vectorized stores. Clone if needed + // to ensure alignment; the result is copied back into self afterwards. + constexpr int vec_bytes = elems_per_call * sizeof(scalar_t); + auto output = self.contiguous(); + if (reinterpret_cast(output.data_ptr()) % vec_bytes != 0) { + output = output.clone(); + } + + constexpr int block_size = 256; + + if (key.dim() == 1) { + // === Launch single key kernel === + constexpr int epc = elems_per_call; + int64_t num_chunks = (self.numel() + epc - 1) / epc; + int num_blocks = static_cast((num_chunks + block_size - 1) / block_size); + + auto key_contig = key.contiguous(); + philox_single_key_kernel + <<>>( + output.mutable_data_ptr(), + key_contig.data_ptr(), + self.numel(), sample_func, param_func); + } else { + // === Launch batched (multiple) key kernel === + // The kernel writes each key's output as a contiguous block of + // elems_per_key elements. We determine elems_per_key by counting + // trailing size-1 key dims; these are the output dimensions that a + // single key generates over. For example, with key shape (4, 1, 1, 2) + // and output shape (4, 10, 100): key_dims=1, elems_per_key=1000. + int64_t elems_per_key = 1; + int64_t key_dims = self.dim(); + for (int64_t i = self.dim() - 1; i >= 0; i--) { + if (key.size(i) != 1) break; + elems_per_key *= self.size(i); + key_dims--; + } + int64_t num_keys = self.numel() / elems_per_key; + + // Handle key, self broadcasting via OffsetCalculator. + c10::SmallVector oc_sizes(key_dims); + c10::SmallVector oc_strides(key_dims); + for (int64_t i = 0; i < key_dims; i++) { + int64_t dim = key_dims - 1 - i; + oc_sizes[i] = self.size(dim); + oc_strides[i] = key.size(dim) > 1 ? key.stride(dim) : 0; + } + const int64_t* oc_strides_ptr = oc_strides.data(); + auto key_offset_calc = OffsetCalculator<1>( + key_dims, oc_sizes.data(), &oc_strides_ptr); + + int64_t chunks_per_key = + (elems_per_key + elems_per_call - 1) / elems_per_call; + int64_t total_threads = num_keys * chunks_per_key; + int num_blocks = static_cast((total_threads + block_size - 1) / block_size); + + philox_multi_key_kernel + <<>>( + output.mutable_data_ptr(), + key.data_ptr(), + num_keys, elems_per_key, + sample_func, param_func, key_offset_calc); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + if (output.data_ptr() != self.data_ptr()) { + self.copy_(output); + } +} + +} // anonymous namespace + +Tensor& _philox_uniform_cuda_( + Tensor& self, const Tensor& key, double low, double high) { + AT_DISPATCH_FLOATING_TYPES_AND2( + kHalf, kBFloat16, self.scalar_type(), "_philox_uniform_", [&] { + auto sample_func = []() { + if constexpr (std::is_same_v) { + return [] __device__ (uint64_t seed, uint64_t offset) { + uint4 r = philox_4x32(seed, offset); + ulonglong2 packed; + packed.x = (static_cast(r.x) << 32) | r.y; + packed.y = (static_cast(r.z) << 32) | r.w; + return packed; + }; + } else { + return [] __device__ (uint64_t seed, uint64_t offset) { + return philox_4x32(seed, offset); + }; + } + }(); + + auto lo = static_cast(low); + auto hi = static_cast(high); + auto param_func = [lo, hi] __device__ (auto rand) { + return static_cast( + at::transformation::uniform_real(rand, lo, hi)); + }; + + philox_distribution_kernel( + "_philox_uniform_", self, key, sample_func, param_func); + }); + return self; +} + +Tensor& _philox_normal_cuda_( + Tensor& self, const Tensor& key, double mean, double stddev) { + AT_DISPATCH_FLOATING_TYPES_AND2( + kHalf, kBFloat16, self.scalar_type(), "_philox_normal_", [&] { + using compute_t = std::conditional_t, double, float>; + auto sample_func = []() { + if constexpr (std::is_same_v) { + return [] __device__ (uint64_t seed, uint64_t offset) { + return box_muller_double(philox_4x32(seed, offset)); + }; + } else { + return [] __device__ (uint64_t seed, uint64_t offset) { + return box_muller_float(philox_4x32(seed, offset)); + }; + } + }(); + + auto mu = static_cast(mean); + auto sigma = static_cast(stddev); + auto param_func = [mu, sigma] __device__ (compute_t rand) { + return static_cast(rand * sigma + mu); + }; + + philox_distribution_kernel( + "_philox_normal_", self, key, sample_func, param_func); + }); + return self; +} + +} // namespace at::native diff --git a/aten/src/ATen/native/cuda/PhiloxKeySplit.cu b/aten/src/ATen/native/cuda/PhiloxKeySplit.cu new file mode 100644 index 0000000000000..07893c061ee0d --- /dev/null +++ b/aten/src/ATen/native/cuda/PhiloxKeySplit.cu @@ -0,0 +1,140 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS + +#include +#include +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#include +#else +#include +#include +#include +#include +#endif + +namespace at::native { + +namespace { + +using at::cuda::philox_4x32; + +// Derive a new (seed, offset) key from 4 random uint32 values. +// Use 2 uint32s for the 64-bit seed and 2 for the 64-bit offset. +__device__ __forceinline__ void philox_derive_key( + uint4 r, + uint64_t* out_seed, + uint64_t* out_offset) { + *out_seed = static_cast(r.x) | (static_cast(r.y) << 32); + *out_offset = static_cast(r.z) | (static_cast(r.w) << 32); +} + +// Grid-stride loop over (split_idx, key_idx) pairs. +__global__ void philox_key_split_kernel( + const uint64_t* __restrict__ input, + uint64_t* __restrict__ output, + int64_t num_keys, + int64_t num_splits) { + int64_t total = num_keys * num_splits; + int64_t tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + for (; tid < total; tid += static_cast(gridDim.x) * blockDim.x) { + int64_t split_idx = tid / num_keys; + int64_t key_idx = tid % num_keys; + + uint64_t seed = input[key_idx * 2]; + uint64_t offset = input[key_idx * 2 + 1]; + + // Sample randomness to get the next (seed, offset pair). + uint4 r = philox_4x32(seed, offset + static_cast(split_idx)); + int64_t out = (split_idx * num_keys + key_idx) * 2; + philox_derive_key(r, &output[out], &output[out + 1]); + } +} + +__global__ void philox_key_fold_in_kernel( + const uint64_t* __restrict__ input, + uint64_t* __restrict__ output, + int64_t num_keys, + int64_t data) { + int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + for (; idx < num_keys; idx += static_cast(gridDim.x) * blockDim.x) { + uint64_t seed = input[idx * 2]; + uint64_t offset = input[idx * 2 + 1]; + + // Sample randomness to get the next (seed, offset pair). + uint4 r = philox_4x32(seed, offset + static_cast(data)); + philox_derive_key(r, &output[idx * 2], &output[idx * 2 + 1]); + } +} + +} // anonymous namespace + +Tensor _philox_key_split_cuda(const Tensor& key, int64_t num_splits) { + TORCH_CHECK(key.dim() >= 1 && key.size(-1) == 2, + "_philox_key_split: key must have shape (*batch, 2), got shape ", + key.sizes()); + TORCH_CHECK(key.scalar_type() == kUInt64, + "_philox_key_split: key must have dtype uint64, got ", + key.scalar_type()); + TORCH_CHECK(num_splits > 0, + "_philox_key_split: num_splits must be positive, got ", + num_splits); + + // Output shape: (num_splits, *key.shape) + auto output_sizes = key.sizes().vec(); + output_sizes.insert(output_sizes.begin(), num_splits); + Tensor output = at::empty(output_sizes, key.options()); + int64_t num_keys = key.numel() / 2; + if (num_keys == 0) { + return output; + } + + int64_t total_threads = num_keys * num_splits; + constexpr int block_size = 256; + int num_blocks = std::min( + static_cast((total_threads + block_size - 1) / block_size), + at::cuda::getCurrentDeviceProperties()->multiProcessorCount * 4); + + auto key_contig = key.contiguous(); + philox_key_split_kernel<<>>( + key_contig.data_ptr(), + output.data_ptr(), + num_keys, num_splits); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return output; +} + +Tensor _philox_key_fold_in_cuda(const Tensor& key, int64_t data) { + TORCH_CHECK(key.dim() >= 1 && key.size(-1) == 2, + "_philox_key_fold_in: key must have shape (*batch, 2), got shape ", + key.sizes()); + TORCH_CHECK(key.scalar_type() == kUInt64, + "_philox_key_fold_in: key must have dtype uint64, got ", + key.scalar_type()); + + Tensor output = at::empty_like(key); + int64_t num_keys = key.numel() / 2; + if (num_keys == 0) { + return output; + } + + constexpr int block_size = 256; + int num_blocks = std::min( + static_cast((num_keys + block_size - 1) / block_size), + at::cuda::getCurrentDeviceProperties()->multiProcessorCount * 4); + + auto key_contig = key.contiguous(); + philox_key_fold_in_kernel<<>>( + key_contig.data_ptr(), + output.data_ptr(), + num_keys, data); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return output; +} + +} // namespace at::native diff --git a/aten/src/ATen/native/cuda/Reduce.cuh b/aten/src/ATen/native/cuda/Reduce.cuh index ee3a05c854540..71084941d5e1d 100644 --- a/aten/src/ATen/native/cuda/Reduce.cuh +++ b/aten/src/ATen/native/cuda/Reduce.cuh @@ -635,10 +635,10 @@ struct ReduceOp { using args_vec_t = std::array; int dim_x = blockDim.x; args_vec_t* shared = (args_vec_t*)shared_memory; - if (dim_x > warpSize) { + if (dim_x > C10_WARP_SIZE) { int address_base = threadIdx.x + threadIdx.y*blockDim.x; shared[address_base] = value; - for (int offset = dim_x/2; offset >= warpSize; offset >>= 1) { + for (int offset = dim_x/2; offset >= C10_WARP_SIZE; offset >>= 1) { __syncthreads(); if (threadIdx.x < offset && threadIdx.x + offset < blockDim.x) { args_vec_t other = shared[address_base + offset]; @@ -649,7 +649,7 @@ struct ReduceOp { shared[address_base] = value; } } - dim_x = warpSize; + dim_x = C10_WARP_SIZE; } __syncthreads(); @@ -836,9 +836,9 @@ struct ReduceOp { } } } else { +#if defined(USE_ROCM) && ROCM_VERSION <= 71300 index_t input_offset = threadIdx.y; index_t step = blockDim.y; -#ifdef USE_ROCM // Prefetch loads to better hide their latency #define PRFCH 4 for (; input_offset < config.ctas_per_output; input_offset += step*PRFCH) { arg_vec_t next[PRFCH]; @@ -855,6 +855,14 @@ struct ReduceOp { } } #else +#if defined(USE_ROCM) + int input_offset = threadIdx.y; + int step = blockDim.y; + #pragma unroll +#else + index_t input_offset = threadIdx.y; + index_t step = blockDim.y; +#endif for (; input_offset < config.ctas_per_output; input_offset += step) { index_t idx = config.staging_memory_offset(input_offset); arg_vec_t next = reduce_buffer[idx]; @@ -1117,7 +1125,7 @@ ReduceConfig setReduceConfig(const TensorIterator& iter){ if (iter.ndim() == 0 || reduction_on_fastest_striding_dimension) { // Split the input across lanes if the input is contiguous in the reduced // dimension. This will require reduction between threads using warp - // shuffle instructions and shared memory (if block_width > warpSize). + // shuffle instructions and shared memory (if block_width > C10_WARP_SIZE). config.input_mult[0] = config.split_input(block_width); } else { // Otherwise split the output across lanes in a warp. diff --git a/aten/src/ATen/native/cuda/RowwiseScaledMM.cu b/aten/src/ATen/native/cuda/RowwiseScaledMM.cu index 960386ac9a431..4bafe771c7e61 100644 --- a/aten/src/ATen/native/cuda/RowwiseScaledMM.cu +++ b/aten/src/ATen/native/cuda/RowwiseScaledMM.cu @@ -198,7 +198,7 @@ void f8f8bf16_rowwise_impl( cutlass::epilogue::collective::EpilogueTileAuto, DtypeAccum, DtypeEpilogue, - DtypeOutput, + void, // Indicate there is no beta scaling to save register LayoutOutput, AlignmentOutput, DtypeOutput, @@ -255,7 +255,7 @@ void f8f8bf16_rowwise_impl( : nullptr}, {{reinterpret_cast(w_scale.data_ptr())}, {{reinterpret_cast(x_scale.data_ptr())}}}}}, - reinterpret_cast(out.data_ptr()), + nullptr, stride_output, reinterpret_cast(out.data_ptr()), stride_output}}; @@ -390,7 +390,7 @@ void f8f8bf16_rowwise_impl_sm100_sm120( TileShape, ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto, DtypeAccum, DtypeEpilogue, - DtypeOutput, LayoutOutput, AlignmentOutput, + void, LayoutOutput, AlignmentOutput, DtypeOutput, LayoutOutput, AlignmentOutput, EpilogueScheduleType, EpilogueEVT>::CollectiveOp; @@ -448,7 +448,7 @@ void f8f8bf16_rowwise_impl_sm100_sm120( : nullptr}, {{reinterpret_cast(w_scale.data_ptr())}, {{reinterpret_cast(x_scale.data_ptr())}}}}}, - reinterpret_cast(out.data_ptr()), + nullptr, stride_output, reinterpret_cast(out.data_ptr()), stride_output}}; diff --git a/aten/src/ATen/native/cuda/ScaledBlas.cpp b/aten/src/ATen/native/cuda/ScaledBlas.cpp index 2c28dc14a0c8a..42d1ea61b17cc 100644 --- a/aten/src/ATen/native/cuda/ScaledBlas.cpp +++ b/aten/src/ATen/native/cuda/ScaledBlas.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include #include #include @@ -826,7 +827,7 @@ _scaled_block1x128_block1x128( scale_a.stride(0) == 1 && ( scale_a.stride(1) == M || - (scale_a.size(1) == 1 && scale_b.stride(1) == 1) + (scale_a.size(1) == 1 && scale_a.stride(1) == 1) ), "scale_a strides must be (", 1, ", ", M, "); got: ", scale_a.strides() ); @@ -848,7 +849,7 @@ _scaled_block1x128_block1x128( scale_b.stride(1) == 1 ) ), - "scale_b strides must be (", 1, ", ", N, "); got: ", scale_a.strides() + "scale_b strides must be (", 1, ", ", N, "); got: ", scale_b.strides() ); auto scaling_choice_a = ScalingType::BlockWise1x128; @@ -986,7 +987,7 @@ _scaled_block1x128_block128x128( scale_a.stride(1) == 1 ) ), - "scale_a must have strides (1, ", M, "); got ", scale_b.strides() + "scale_a must have strides (1, ", M, "); got ", scale_a.strides() ); // scale_b shape TORCH_CHECK_VALUE( @@ -1228,7 +1229,7 @@ _scaled_nvfp4_nvfp4( void check_swizzle_lengths(ScaledGemmImplementation impl, std::vector& swizzle_a, std::vector& swizzle_b) { -#ifdef ROCM +#ifdef USE_ROCM // ROCM doesn't swizzle their formats - we don't care what's passed. return; #else diff --git a/aten/src/ATen/native/cuda/ScaledGroupMM.cu b/aten/src/ATen/native/cuda/ScaledGroupMM.cu index a276579b15fd2..69a1a27c76bb1 100644 --- a/aten/src/ATen/native/cuda/ScaledGroupMM.cu +++ b/aten/src/ATen/native/cuda/ScaledGroupMM.cu @@ -192,7 +192,7 @@ void f8f8bf16_grouped_gemm_impl_sm90( cutlass::epilogue::collective::EpilogueTileAuto, DtypeAccum, DtypeAccum, - DtypeOutput, + void, // Indicate there is no beta scaling to save register LayoutOutput*, AlignmentOutput, DtypeOutput, @@ -390,7 +390,7 @@ void f8f8bf16_grouped_gemm_impl_sm90( (const DtypeB**)inputB_ptrs, stride_B}, {{{{inputB_scale_ptrs}, {{inputA_scale_ptrs}, {}, {}}, {}}, {}}, - (const DtypeOutput**)output_ptrs, + nullptr, stride_output, output_ptrs, stride_output}}; diff --git a/aten/src/ATen/native/cuda/Sort.cu b/aten/src/ATen/native/cuda/Sort.cu index bae17a6a01236..a05ded2e2bb0f 100644 --- a/aten/src/ATen/native/cuda/Sort.cu +++ b/aten/src/ATen/native/cuda/Sort.cu @@ -107,7 +107,7 @@ struct SmallBitonicSort { // For small sorts (n <= 128) we use warpMergeSortKVInPlace which // sorts one slice per warp and potentially multiple slices in the // same block for improved occupancy with large batch sizes. -template +template struct WarpMergeSort { template @@ -128,7 +128,7 @@ struct WarpMergeSort { const auto min_grid = minimum_grid_for_occupancy( warpMergeSortKVInPlace< A, -1, sort_size, max_block_y, - K, V, LTOp, IndexType>, + K, V, LTOp, IndexType, WARP_SIZE>, block_x * max_block_y); const auto max_batch = std::max(IndexType{1}, keySlices / min_grid); const int block_y = std::min(IndexType(max_block_y), max_batch); @@ -142,7 +142,9 @@ struct WarpMergeSort { if (descending) { const K invalid_key = at::numeric_limits::lower_bound(); - warpMergeSortKVInPlace + warpMergeSortKVInPlace< + A, -1, sort_size, max_block_y, + K, V, GTOp, IndexType, WARP_SIZE> <<>>( keyInfo, keySlices, @@ -161,7 +163,9 @@ struct WarpMergeSort { } return at::numeric_limits::upper_bound(); }(); - warpMergeSortKVInPlace + warpMergeSortKVInPlace< + A, -1, sort_size, max_block_y, + K, V, LTOp, IndexType, WARP_SIZE> <<>>( keyInfo, keySlices, @@ -375,7 +379,16 @@ void sortKeyValueInplace( sortCommon(SmallBitonicSort{}, key, value, dim, descending); #if HAS_WARP_MERGE_SORT() } else if (sort_size <= 128) { - sortCommon(WarpMergeSort<128>{}, key, value, dim, descending); +#ifdef USE_ROCM + if (at::cuda::warp_size() == 32) { + sortCommon(WarpMergeSort<128, 32>{}, key, value, dim, descending); + } + else { + sortCommon(WarpMergeSort<128, 64>{}, key, value, dim, descending); + } +#else + sortCommon(WarpMergeSort<128, C10_WARP_SIZE>{}, key, value, dim, descending); +#endif #endif } else { sortCommon(MediumRadixSort{}, key, value, dim, descending); diff --git a/aten/src/ATen/native/cuda/SortUtils.cuh b/aten/src/ATen/native/cuda/SortUtils.cuh index 8e424707c7ef6..8bcb0ae323cfc 100644 --- a/aten/src/ATen/native/cuda/SortUtils.cuh +++ b/aten/src/ATen/native/cuda/SortUtils.cuh @@ -165,15 +165,18 @@ bitonicSortKVInPlace(at::cuda::detail::TensorInfo keys, #if HAS_WARP_MERGE_SORT() +// Note [warp merge sort WARP_SIZE template param] +// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +// warpMergeSortKVInPlace was written assuming C10_WARP_SIZE is a constexpr. +// In torch/headeronly/macros/Macros.h, C10_WARP_SIZE is 32 for CUDA, and on +// ROCm it will be 32 or 64 based on the current compile-time gfx target. +// Ideally, warpSize should be used instead of C10_WARP_SIZE in device code, but +// C10_WARP_SIZE within this kernel has been used as a template parameter for +// some device functions. Therefore, a template param for WARP_SIZE was added. template -#if !defined(USE_ROCM) -// On CUDA, use explicit launch bounds for better occupancy -C10_LAUNCH_BOUNDS_1(C10_WARP_SIZE * max_block_dim_y) -#endif -// Note: ROCm doesn't use launch bounds here because C10_WARP_SIZE is not -// a true compile-time constant in device code (it's a constexpr function). -// The compiler infers good launch bounds from the kernel code automatically. + typename K, typename V, typename Comparator, typename IndexType, + int WARP_SIZE> +C10_LAUNCH_BOUNDS_1(WARP_SIZE * max_block_dim_y) __global__ void warpMergeSortKVInPlace( at::cuda::detail::TensorInfo keys, @@ -207,17 +210,17 @@ warpMergeSortKVInPlace( namespace cub = ROCM_HIPCUB(at_cuda_detail::cub); - CUDA_KERNEL_ASSERT(blockDim.x == C10_WARP_SIZE); + CUDA_KERNEL_ASSERT(blockDim.x == WARP_SIZE); CUDA_KERNEL_ASSERT(blockDim.y <= max_block_dim_y); - constexpr int items_per_thread = sort_size / C10_WARP_SIZE; + constexpr int items_per_thread = sort_size / WARP_SIZE; static_assert( - items_per_thread * C10_WARP_SIZE == sort_size, - "sort_size must be a multiple of C10_WARP_SIZE"); + items_per_thread * WARP_SIZE == sort_size, + "sort_size must be a multiple of WARP_SIZE template param"); using LoadKeys = cub::WarpLoad; using LoadValues = cub::WarpLoad; - using Sort = cub::WarpMergeSort; + using Sort = cub::WarpMergeSort; using StoreKeys = cub::WarpStore; using StoreValues = cub::WarpStore; diff --git a/aten/src/ATen/native/cuda/SortingRadixSelect.cuh b/aten/src/ATen/native/cuda/SortingRadixSelect.cuh index 4ebc6de07e534..67c15e0dd3385 100644 --- a/aten/src/ATen/native/cuda/SortingRadixSelect.cuh +++ b/aten/src/ATen/native/cuda/SortingRadixSelect.cuh @@ -474,13 +474,13 @@ __device__ __forceinline__ void countRadixAggregateCounts( // Maximum number of warps per workgroup. HIP workgroups have at most 1024 threads. // Warp size is at least 32 (can be 64 on some architectures), so we use 32 for safety. // This sizes shared memory buffers to accommodate all possible warps: 1024/32 = 32. - constexpr uint MAX_WARPS = 1024/32; + constexpr uint MAX_WARPS = 1024/C10_WARP_SIZE_LOWER_BOUND; const int buffer_offset = buffer_index * MAX_WARPS * RadixSize; // offset of the buffer in smem. - const uint WARP_BITS = __builtin_ctz(warpSize); + const uint WARP_BITS = __builtin_ctz(C10_WARP_SIZE); const uint num_warps = blockDim.x >> WARP_BITS; // Actual number of warps in this block - const uint warp_id = threadIdx.x >> WARP_BITS; // = threadIdx.x / warpSize - const int lane_id = at::cuda::getLaneId(); // = threadIdx.x % warpSize + const uint warp_id = threadIdx.x >> WARP_BITS; // = threadIdx.x / C10_WARP_SIZE + const int lane_id = at::cuda::getLaneId(); // = threadIdx.x % C10_WARP_SIZE // Stage 1: Each warp's lane 0 stores its counts in smem. // Layout after Stage 1: [warp0: all radix bins], [warp1: all radix bins], ... @@ -521,7 +521,6 @@ __device__ __forceinline__ void countRadixAggregateCounts( for (uint32_t i = 0; i < RadixSize; ++i) { counts[i] = smem[buffer_offset + i]; } - __syncthreads(); // Wait for all threads to finish reading the final counts. } // This function counts the distribution of all input values in a @@ -692,6 +691,15 @@ __device__ scalar_t findPatternDataSmem( const scalar_t* dataSmem, // input data stored in shared memory. index_t dataSmemSize) { // input data size stored in shared memory. + // Ensure all threads have finished reading from smem before overwriting it. + // countRadixAggregateCounts Stage 3 reads from smem[buffer_offset + i]; + // when buffer_offset == 0, those locations overlap with smem[0]/smem[1] + // written below. Warp 0 (which writes smem[0]/smem[1]) may get ahead of + // lagging warps still in Stage 3. Syncing here (rather than at the end of + // Stage 3) is cheaper because findPatternDataSmem is called at most once per + // radixSelect invocation, only when a unique element is found (count == 1). + __syncthreads(); + // initialize smem to 0. // smem[0] is a flag to indicate if a value has been found. // smem[1] is the found value. diff --git a/aten/src/ATen/native/cuda/TensorTopK.cu b/aten/src/ATen/native/cuda/TensorTopK.cu index 835b718fb6d49..2d21944e92ceb 100644 --- a/aten/src/ATen/native/cuda/TensorTopK.cu +++ b/aten/src/ATen/native/cuda/TensorTopK.cu @@ -319,7 +319,7 @@ __global__ void gatherTopK(at::cuda::detail::TensorInfo inpu } __syncthreads(); // All threads within the warp need to participate in the loop, so rounding up to a multiple of the warp size. - IndexType numIterations = round_up(inputSliceSize, (IndexType) warpSize); + IndexType numIterations = round_up(inputSliceSize, (IndexType) C10_WARP_SIZE); // phase 1: write actual > `pattern` (or < `pattern`, depending on the sort direction) values to the output. // prefetching data from global memory. @@ -447,8 +447,9 @@ constexpr int MAX_WARP_TOPK_SLICE = 512; // GTOp/LTOp instead of bitwise conversion. Bitwise conversion is only needed for radix sorting. // Kernel using WarpMergeSort for small topK operations +// See Note [warp merge sort WARP_SIZE template param] template + typename scalar_t, typename IndexType, bool is_descending, int WARP_SIZE> __global__ void warpMergeSortTopK( at::cuda::detail::TensorInfo input, IndexType inputSliceSize, @@ -486,14 +487,14 @@ __global__ void warpMergeSortTopK( namespace cub = ROCM_HIPCUB(at_cuda_detail::cub); - CUDA_KERNEL_ASSERT(blockDim.x == C10_WARP_SIZE); + CUDA_KERNEL_ASSERT(blockDim.x == WARP_SIZE); CUDA_KERNEL_ASSERT(blockDim.y <= max_block_dim_y); - constexpr int items_per_thread = sort_size / C10_WARP_SIZE; - static_assert(items_per_thread * C10_WARP_SIZE == sort_size, - "sort_size must be a multiple of C10_WARP_SIZE"); + constexpr int items_per_thread = sort_size / WARP_SIZE; + static_assert(items_per_thread * WARP_SIZE == sort_size, + "sort_size must be a multiple of WARP_SIZE template param"); using LoadKeys = cub::WarpLoad; - using Sort = cub::WarpMergeSort; + using Sort = cub::WarpMergeSort; using StoreKeys = cub::WarpStore; using StoreIndices = cub::WarpStore; @@ -569,12 +570,30 @@ void launch( "Too many slices for warp topk"); // Dispatch based on sort size and sort direction + // See Note [warp merge sort WARP_SIZE template param] +#ifdef USE_ROCM #define LAUNCH_KERNEL(SORT_SIZE, IS_DESCENDING) \ - warpMergeSortTopK \ + if (at::cuda::warp_size() == 32) { \ + warpMergeSortTopK \ + <<>>( \ + input, inputSliceSize, k, numInputSlices, inputWithinSliceStride, \ + topK, topKWithinSliceStride, indices, indicesWithinSliceStride); \ + } \ + else { \ + warpMergeSortTopK \ + <<>>( \ + input, inputSliceSize, k, numInputSlices, inputWithinSliceStride, \ + topK, topKWithinSliceStride, indices, indicesWithinSliceStride); \ + } \ + C10_CUDA_KERNEL_LAUNCH_CHECK() +#else + #define LAUNCH_KERNEL(SORT_SIZE, IS_DESCENDING) \ + warpMergeSortTopK \ <<>>( \ input, inputSliceSize, k, numInputSlices, inputWithinSliceStride, \ topK, topKWithinSliceStride, indices, indicesWithinSliceStride); \ C10_CUDA_KERNEL_LAUNCH_CHECK() +#endif // We have specialized launches for different sizes, as sort_size affects // shared memory, registers per thread and occupancy. We can use 'LAUNCH_KERNEL(512, false);' @@ -849,8 +868,9 @@ __global__ void computeBlockwiseWithinKCounts( } } - constexpr int num_warps = RADIX_DIGITS / C10_WARP_SIZE; - __shared__ uint32_t warp_counts[num_warps]; + constexpr int SHMEM_SIZE = RADIX_DIGITS / C10_WARP_SIZE_LOWER_BOUND; // max shmem size on ROCm + const int num_warps = RADIX_DIGITS / C10_WARP_SIZE; + __shared__ uint32_t warp_counts[SHMEM_SIZE]; if (tidx % C10_WARP_SIZE == 0) { warp_counts[warp] = count; } @@ -1023,8 +1043,9 @@ __global__ void computeBlockwiseWithinKCounts( } } - constexpr int num_warps = RADIX_DIGITS / C10_WARP_SIZE; - __shared__ uint32_t warp_counts[num_warps]; + constexpr int SHMEM_SIZE = RADIX_DIGITS / C10_WARP_SIZE_LOWER_BOUND; // max shmem size on ROCm + const int num_warps = RADIX_DIGITS / C10_WARP_SIZE; + __shared__ uint32_t warp_counts[SHMEM_SIZE]; if (tidx % C10_WARP_SIZE == 0) { warp_counts[warp] = count; } diff --git a/aten/src/ATen/native/cuda/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h b/aten/src/ATen/native/cuda/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h index dc07a53543ad7..f75606bdc9eae 100644 --- a/aten/src/ATen/native/cuda/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h +++ b/aten/src/ATen/native/cuda/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h @@ -68,13 +68,13 @@ struct GemmFpAIntB { using ElementA = typename Mma::IteratorA::Element; using LayoutA = typename Mma::IteratorA::Layout; using ElementB = typename Mma::IteratorB::Element; - using LayoutB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; using ElementC = typename Epilogue::OutputTileIterator::Element; using LayoutC = typename Mma::LayoutC; using ElementScale = ElementC; static ComplexTransform const kTransformA = Mma::kTransformA; - static ComplexTransform const kTransformB = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; // Type definitions about the mainloop. using Operator = typename Mma::Operator; diff --git a/aten/src/ATen/native/cuda/fused_adagrad_utils.cuh b/aten/src/ATen/native/cuda/fused_adagrad_utils.cuh index cb82c30fae198..fd70a33affdcf 100644 --- a/aten/src/ATen/native/cuda/fused_adagrad_utils.cuh +++ b/aten/src/ATen/native/cuda/fused_adagrad_utils.cuh @@ -104,6 +104,9 @@ struct FusedAdagradMathFunctor { found_inf_ptr); load_store(args[kParamIdx], r_args[kParamIdx], i_start, 0); + if (grad_scale_ptr) { + load_store(args[kGradIdx], r_args[kGradIdx], i_start, 0); + } load_store(args[kStateSumIdx], r_args[kStateSumIdx], i_start, 0); } } else { diff --git a/aten/src/ATen/native/cuda/fused_adam_amsgrad_impl.cu b/aten/src/ATen/native/cuda/fused_adam_amsgrad_impl.cu index cef07de1b41f9..f4844a79b2d70 100644 --- a/aten/src/ATen/native/cuda/fused_adam_amsgrad_impl.cu +++ b/aten/src/ATen/native/cuda/fused_adam_amsgrad_impl.cu @@ -1,6 +1,7 @@ #include #include +#include #include #include #include @@ -36,26 +37,64 @@ void _fused_adam_amsgrad_cuda_impl_( found_inf.has_value() ? found_inf->data_ptr() : nullptr; const float* lr_ptr = nullptr; - AT_DISPATCH_FLOATING_TYPES_AND2( - kHalf, - kBFloat16, - params[0].scalar_type(), - "fused_adam_kernel_cuda", - [&]() { - multi_tensor_apply_for_fused_optimizer<5>( - tensor_lists, - state_steps, - FusedAdamMathFunctor(), - lr_ptr, // unused - lr, - beta1, - beta2, - weight_decay, - eps, - maximize, - grad_scale_ptr, - found_inf_ptr); - }); + if (params[0].scalar_type() != exp_avgs[0].scalar_type()) { + validate_mixed_precision_dtypes( + params, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + "Mixed-precision fused Adam"); + AT_DISPATCH_V2( + exp_avgs[0].scalar_type(), + "fused_adam_amsgrad_mp_kernel_cuda", + AT_WRAP([&]() { + multi_tensor_apply_for_fused_optimizer<5>( + tensor_lists, + state_steps, + FusedAdamMathFunctorMP< + float, + float, + float, + scalar_t, + scalar_t, + scalar_t, + 5, + ADAM_MODE::ORIGINAL, + true>(), + lr_ptr, // unused + lr, + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }), + kBFloat16); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + kHalf, + kBFloat16, + params[0].scalar_type(), + "fused_adam_kernel_cuda", + [&]() { + multi_tensor_apply_for_fused_optimizer<5>( + tensor_lists, + state_steps, + FusedAdamMathFunctor(), + lr_ptr, // unused + lr, + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }); + } } // The following overload simply has a Tensor lr @@ -87,26 +126,64 @@ void _fused_adam_amsgrad_cuda_impl_( found_inf.has_value() ? found_inf->data_ptr() : nullptr; const float* lr_ptr = lr.const_data_ptr(); - AT_DISPATCH_FLOATING_TYPES_AND2( - kHalf, - kBFloat16, - params[0].scalar_type(), - "fused_adam_kernel_cuda", - [&]() { - multi_tensor_apply_for_fused_optimizer<5>( - tensor_lists, - state_steps, - FusedAdamMathFunctor(), - lr_ptr, - 1.0, // unused - beta1, - beta2, - weight_decay, - eps, - maximize, - grad_scale_ptr, - found_inf_ptr); - }); + if (params[0].scalar_type() != exp_avgs[0].scalar_type()) { + validate_mixed_precision_dtypes( + params, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + "Mixed-precision fused Adam"); + AT_DISPATCH_V2( + exp_avgs[0].scalar_type(), + "fused_adam_amsgrad_mp_kernel_cuda", + AT_WRAP([&]() { + multi_tensor_apply_for_fused_optimizer<5>( + tensor_lists, + state_steps, + FusedAdamMathFunctorMP< + float, + float, + float, + scalar_t, + scalar_t, + scalar_t, + 5, + ADAM_MODE::ORIGINAL, + true>(), + lr_ptr, + 1.0, // unused + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }), + kBFloat16); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + kHalf, + kBFloat16, + params[0].scalar_type(), + "fused_adam_kernel_cuda", + [&]() { + multi_tensor_apply_for_fused_optimizer<5>( + tensor_lists, + state_steps, + FusedAdamMathFunctor(), + lr_ptr, + 1.0, // unused + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }); + } } } // namespace at::native diff --git a/aten/src/ATen/native/cuda/fused_adam_impl.cu b/aten/src/ATen/native/cuda/fused_adam_impl.cu index 2c1f5ce0d6d57..e51e446ea5e74 100644 --- a/aten/src/ATen/native/cuda/fused_adam_impl.cu +++ b/aten/src/ATen/native/cuda/fused_adam_impl.cu @@ -1,6 +1,7 @@ #include #include +#include #include #include #include @@ -31,26 +32,59 @@ void _fused_adam_cuda_impl_( found_inf.has_value() ? found_inf->data_ptr() : nullptr; const float* lr_ptr = nullptr; - AT_DISPATCH_FLOATING_TYPES_AND2( - kHalf, - kBFloat16, - params[0].scalar_type(), - "fused_adam_kernel_cuda", - [&]() { - multi_tensor_apply_for_fused_optimizer<4>( - tensor_lists, - state_steps, - FusedAdamMathFunctor(), - lr_ptr, // unused - lr, - beta1, - beta2, - weight_decay, - eps, - maximize, - grad_scale_ptr, - found_inf_ptr); - }); + if (params[0].scalar_type() != exp_avgs[0].scalar_type()) { + validate_mixed_precision_dtypes( + params, grads, exp_avgs, exp_avg_sqs, "Mixed-precision fused Adam"); + AT_DISPATCH_V2( + exp_avgs[0].scalar_type(), + "fused_adam_mp_kernel_cuda", + AT_WRAP([&]() { + multi_tensor_apply_for_fused_optimizer<4>( + tensor_lists, + state_steps, + FusedAdamMathFunctorMP< + float, + float, + float, + scalar_t, + scalar_t, + float, + 4, + ADAM_MODE::ORIGINAL, + false>(), + lr_ptr, // unused + lr, + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }), + kBFloat16); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + kHalf, + kBFloat16, + params[0].scalar_type(), + "fused_adam_kernel_cuda", + [&]() { + multi_tensor_apply_for_fused_optimizer<4>( + tensor_lists, + state_steps, + FusedAdamMathFunctor(), + lr_ptr, // unused + lr, + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }); + } } // The following overload simply has a Tensor lr @@ -77,26 +111,59 @@ void _fused_adam_cuda_impl_( found_inf.has_value() ? found_inf->data_ptr() : nullptr; const float* lr_ptr = lr.const_data_ptr(); - AT_DISPATCH_FLOATING_TYPES_AND2( - kHalf, - kBFloat16, - params[0].scalar_type(), - "fused_adam_kernel_cuda", - [&]() { - multi_tensor_apply_for_fused_optimizer<4>( - tensor_lists, - state_steps, - FusedAdamMathFunctor(), - lr_ptr, - 1.0, // unused - beta1, - beta2, - weight_decay, - eps, - maximize, - grad_scale_ptr, - found_inf_ptr); - }); + if (params[0].scalar_type() != exp_avgs[0].scalar_type()) { + validate_mixed_precision_dtypes( + params, grads, exp_avgs, exp_avg_sqs, "Mixed-precision fused Adam"); + AT_DISPATCH_V2( + exp_avgs[0].scalar_type(), + "fused_adam_mp_kernel_cuda", + AT_WRAP([&]() { + multi_tensor_apply_for_fused_optimizer<4>( + tensor_lists, + state_steps, + FusedAdamMathFunctorMP< + float, + float, + float, + scalar_t, + scalar_t, + float, + 4, + ADAM_MODE::ORIGINAL, + false>(), + lr_ptr, + 1.0, // unused + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }), + kBFloat16); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + kHalf, + kBFloat16, + params[0].scalar_type(), + "fused_adam_kernel_cuda", + [&]() { + multi_tensor_apply_for_fused_optimizer<4>( + tensor_lists, + state_steps, + FusedAdamMathFunctor(), + lr_ptr, + 1.0, // unused + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }); + } } } // namespace at::native diff --git a/aten/src/ATen/native/cuda/fused_adam_utils.cuh b/aten/src/ATen/native/cuda/fused_adam_utils.cuh index 0533db28558b7..5440577ded0ad 100644 --- a/aten/src/ATen/native/cuda/fused_adam_utils.cuh +++ b/aten/src/ATen/native/cuda/fused_adam_utils.cuh @@ -3,12 +3,73 @@ #include #include #include +#include #include namespace at::native { enum class ADAM_MODE : uint8_t { ORIGINAL = 0, ADAMW = 1 }; +// Validates the dtype configuration for mixed-precision fused Adam/AdamW. +// +// Currently the only supported configuration is: +// params/grads: float32 +// optimizer states (exp_avg, exp_avg_sq, ...): bfloat16 +// +// This specific configuration (fp32 params + bf16 optimizer states) has been +// validated end-to-end in large-scale training runs (e.g. DeepSeek-V3 671B) +// and is the only one for which training convergence has been demonstrated. +// Additional mixed-precision configurations (e.g. float16 states) can be +// enabled here once convergence is verified for those as well. +// +// Only [0] is checked because within-list dtype homogeneity is guaranteed by +// _check_tensors_share_device_and_dtype (with skip_cross_list_dtype_check) +// and the Python-side grouping in +// _group_tensors_by_first_tensors_device_and_dtype. +inline void validate_mixed_precision_dtypes( + at::TensorList params, + at::TensorList grads, + at::TensorList exp_avgs, + at::TensorList exp_avg_sqs, + const char* op_name) { + TORCH_CHECK( + params[0].scalar_type() == at::kFloat, + op_name, + " requires float32 params, got ", + params[0].scalar_type()); + TORCH_CHECK( + grads[0].scalar_type() == at::kFloat, + op_name, + " requires float32 grads, got ", + grads[0].scalar_type()); + TORCH_CHECK( + exp_avgs[0].scalar_type() == at::kBFloat16, + op_name, + " requires bfloat16 optimizer states, got ", + exp_avgs[0].scalar_type()); + TORCH_CHECK( + exp_avg_sqs[0].scalar_type() == at::kBFloat16, + op_name, + " requires bfloat16 optimizer states, got ", + exp_avg_sqs[0].scalar_type()); +} + +inline void validate_mixed_precision_dtypes( + at::TensorList params, + at::TensorList grads, + at::TensorList exp_avgs, + at::TensorList exp_avg_sqs, + at::TensorList max_exp_avg_sqs, + const char* op_name) { + validate_mixed_precision_dtypes( + params, grads, exp_avgs, exp_avg_sqs, op_name); + TORCH_CHECK( + max_exp_avg_sqs[0].scalar_type() == at::kBFloat16, + op_name, + " requires bfloat16 max_exp_avg_sqs, got ", + max_exp_avg_sqs[0].scalar_type()); +} + namespace { constexpr uint8_t kParamIdx = 0; @@ -211,6 +272,282 @@ struct FusedAdamMathFunctor { } } }; + +template < + typename scalar_type, + typename param_type, + typename grad_type, + typename exp_avg_type, + typename exp_avg_sq_type, + typename max_exp_avg_sq_type, + int depth, + ADAM_MODE adam_mode, + bool amsgrad> +struct FusedAdamMathFunctorMP { + static_assert( + depth == 4 || depth == 5, + "depth of 4 for Adam, depth of 5 for Adam with AMSGrad."); + using opmath_t = at::opmath_type; + C10_DEVICE __forceinline__ void operator()( + int64_t chunk_size, + FusedOptimizerTensorListMetadata& tl, + const float* lr_ptr, + const double& lr, + const double& beta1, + const double& beta2, + const double& weight_decay, + const double& eps, + const bool& maximize, + const float* grad_scale_ptr, + const float* found_inf_ptr) { + const auto tensor_loc = tl.block_to_tensor[blockIdx.x]; + const auto chunk_idx = tl.block_to_chunk[blockIdx.x]; + const double lr_double = lr_ptr ? *lr_ptr : lr; + + if (found_inf_ptr && *found_inf_ptr == 1) { + return; + } + const auto [bias_correction1, bias_correction2_sqrt] = + [&]() -> std::pair { + auto* step_count = + reinterpret_cast(tl.state_steps_addresses[tensor_loc]); + const auto bias_correction1 = 1 - at::native::pow_(beta1, *step_count); + const auto bias_correction2 = 1 - at::native::pow_(beta2, *step_count); + const auto bias_correction2_sqrt = std::sqrt(bias_correction2); + return {bias_correction1, bias_correction2_sqrt}; + }(); + + param_type* param_args; + grad_type* grad_args; + exp_avg_type* exp_avg_args; + exp_avg_sq_type* exp_avg_sq_args; + [[maybe_unused]] max_exp_avg_sq_type* max_exp_avg_sq_args; + + // r_args represents the state when everything is casted to scalar_type + // to be passed into the adam_math function. scalar_type is our operation + // math type. + scalar_type r_args[depth][kILP]; + + // n = total numel of tensor - what's already been processed + // so n = numel in current tensor not yet processed + const auto n = tl.numel_for_tensor[tensor_loc] - chunk_idx * chunk_size; + + bool all_aligned = init_args_mixed_prec< + depth, + param_type, + grad_type, + exp_avg_type, + exp_avg_sq_type>( + ¶m_args, + &grad_args, + &exp_avg_args, + &exp_avg_sq_args, + tl, + chunk_idx, + chunk_size, + tensor_loc); + if constexpr (amsgrad) { + max_exp_avg_sq_args = + (max_exp_avg_sq_type*)tl.addresses[kMaxExpAvgSqIdx][tensor_loc] + + chunk_idx * chunk_size; + all_aligned = all_aligned && is_aligned(max_exp_avg_sq_args); + } + if ((n % kILP == 0) && (chunk_size % kILP == 0) && all_aligned) { + for (int64_t i_start = threadIdx.x; + i_start * kILP < n && i_start * kILP < chunk_size; + i_start += blockDim.x) { + if constexpr (!std::is_same_v) { + scalar_type casted_param_args[kILP]; + for (int ii = 0; ii < kILP; ii++) { + casted_param_args[ii] = + static_cast(param_args[ii + i_start * kILP]); + } + load_store(r_args[kParamIdx], casted_param_args, 0, 0); + } else { + load_store(r_args[kParamIdx], (scalar_type*)param_args, 0, i_start); + } + if constexpr (!std::is_same_v) { + scalar_type casted_grad_args[kILP]; + for (int ii = 0; ii < kILP; ii++) { + casted_grad_args[ii] = + static_cast(grad_args[ii + i_start * kILP]); + } + load_store(r_args[kGradIdx], casted_grad_args, 0, 0); + } else { + load_store(r_args[kGradIdx], (scalar_type*)grad_args, 0, i_start); + } + if constexpr (!std::is_same_v) { + scalar_type casted_exp_avg_args[kILP]; + for (int ii = 0; ii < kILP; ii++) { + casted_exp_avg_args[ii] = + static_cast(exp_avg_args[ii + i_start * kILP]); + } + load_store(r_args[kExpAvgIdx], casted_exp_avg_args, 0, 0); + } else { + load_store( + r_args[kExpAvgIdx], (scalar_type*)exp_avg_args, 0, i_start); + } + if constexpr (!std::is_same_v) { + scalar_type casted_exp_avg_sq_args[kILP]; + for (int ii = 0; ii < kILP; ii++) { + casted_exp_avg_sq_args[ii] = + static_cast(exp_avg_sq_args[ii + i_start * kILP]); + } + load_store(r_args[kExpAvgSqIdx], casted_exp_avg_sq_args, 0, 0); + } else { + load_store( + r_args[kExpAvgSqIdx], (scalar_type*)exp_avg_sq_args, 0, i_start); + } + if constexpr (amsgrad) { + if constexpr (!std::is_same_v) { + scalar_type casted[kILP]; + for (int ii = 0; ii < kILP; ii++) { + casted[ii] = static_cast( + max_exp_avg_sq_args[ii + i_start * kILP]); + } + load_store(r_args[kMaxExpAvgSqIdx], casted, 0, 0); + } else { + load_store( + r_args[kMaxExpAvgSqIdx], + (scalar_type*)max_exp_avg_sq_args, + 0, + i_start); + } + } + adam_math( + r_args, + lr_double, + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr, + bias_correction1, + bias_correction2_sqrt); + if constexpr (!std::is_same_v) { + param_type casted_r_args[kILP]; + for (int ii = 0; ii < kILP; ii++) { + casted_r_args[ii] = static_cast(r_args[kParamIdx][ii]); + } + load_store(param_args, casted_r_args, i_start, 0); + } else { + load_store(param_args, (param_type*)r_args[kParamIdx], i_start, 0); + } + if constexpr (!std::is_same_v) { + exp_avg_type casted_r_args[kILP]; + for (int ii = 0; ii < kILP; ii++) { + casted_r_args[ii] = + static_cast(r_args[kExpAvgIdx][ii]); + } + load_store(exp_avg_args, casted_r_args, i_start, 0); + } else { + load_store( + exp_avg_args, (exp_avg_type*)r_args[kExpAvgIdx], i_start, 0); + } + if constexpr (!std::is_same_v) { + exp_avg_sq_type casted_r_args[kILP]; + for (int ii = 0; ii < kILP; ii++) { + casted_r_args[ii] = + static_cast(r_args[kExpAvgSqIdx][ii]); + } + load_store(exp_avg_sq_args, casted_r_args, i_start, 0); + } else { + load_store( + exp_avg_sq_args, + (exp_avg_sq_type*)r_args[kExpAvgSqIdx], + i_start, + 0); + } + if constexpr (amsgrad) { + if constexpr (!std::is_same_v) { + max_exp_avg_sq_type casted[kILP]; + for (int ii = 0; ii < kILP; ii++) { + casted[ii] = + static_cast(r_args[kMaxExpAvgSqIdx][ii]); + } + load_store(max_exp_avg_sq_args, casted, i_start, 0); + } else { + load_store( + max_exp_avg_sq_args, + (max_exp_avg_sq_type*)r_args[kMaxExpAvgSqIdx], + i_start, + 0); + } + } + if (grad_scale_ptr) { + if constexpr (!std::is_same_v) { + grad_type casted_r_args[kILP]; + for (int ii = 0; ii < kILP; ii++) { + casted_r_args[ii] = static_cast(r_args[kGradIdx][ii]); + } + load_store(grad_args, casted_r_args, i_start, 0); + } else { + load_store(grad_args, (grad_type*)r_args[kGradIdx], i_start, 0); + } + } + } + } else { + for (int64_t i_start = 0; i_start < n && i_start < chunk_size; + i_start += blockDim.x * kILP) { + load_args< + scalar_type, + param_type, + grad_type, + exp_avg_type, + exp_avg_sq_type>( + r_args, + param_args, + grad_args, + exp_avg_args, + exp_avg_sq_args, + i_start, + chunk_size, + n); + if constexpr (amsgrad) { +#pragma unroll + for (int ii = 0; ii < kILP; ii++) { + const auto i = i_start + threadIdx.x + ii * blockDim.x; + r_args[kMaxExpAvgSqIdx][ii] = 0; + if (i < n && i < chunk_size) { + r_args[kMaxExpAvgSqIdx][ii] = + static_cast(max_exp_avg_sq_args[i]); + } + } + } + adam_math( + r_args, + lr_double, + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr, + bias_correction1, + bias_correction2_sqrt); + store_args(param_args, r_args[kParamIdx], i_start, chunk_size, n); + store_args(exp_avg_args, r_args[kExpAvgIdx], i_start, chunk_size, n); + store_args( + exp_avg_sq_args, r_args[kExpAvgSqIdx], i_start, chunk_size, n); + if constexpr (amsgrad) { + store_args( + max_exp_avg_sq_args, + r_args[kMaxExpAvgSqIdx], + i_start, + chunk_size, + n); + } + if (grad_scale_ptr) { + store_args(grad_args, r_args[kGradIdx], i_start, chunk_size, n); + } + } + } + } +}; + } // namespace } // namespace at::native diff --git a/aten/src/ATen/native/cuda/fused_adamw_amsgrad_impl.cu b/aten/src/ATen/native/cuda/fused_adamw_amsgrad_impl.cu index b2eff4839133f..dba2b2c2af830 100644 --- a/aten/src/ATen/native/cuda/fused_adamw_amsgrad_impl.cu +++ b/aten/src/ATen/native/cuda/fused_adamw_amsgrad_impl.cu @@ -1,6 +1,7 @@ #include #include +#include #include #include #include @@ -36,26 +37,64 @@ void _fused_adamw_amsgrad_cuda_impl_( found_inf.has_value() ? found_inf->data_ptr() : nullptr; const float* lr_ptr = nullptr; - AT_DISPATCH_FLOATING_TYPES_AND2( - kHalf, - kBFloat16, - params[0].scalar_type(), - "fused_adamw_kernel_cuda", - [&]() { - multi_tensor_apply_for_fused_optimizer<5>( - tensor_lists, - state_steps, - FusedAdamMathFunctor(), - lr_ptr, // unused - lr, - beta1, - beta2, - weight_decay, - eps, - maximize, - grad_scale_ptr, - found_inf_ptr); - }); + if (params[0].scalar_type() != exp_avgs[0].scalar_type()) { + validate_mixed_precision_dtypes( + params, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + "Mixed-precision fused AdamW"); + AT_DISPATCH_V2( + exp_avgs[0].scalar_type(), + "fused_adamw_amsgrad_mp_kernel_cuda", + AT_WRAP([&]() { + multi_tensor_apply_for_fused_optimizer<5>( + tensor_lists, + state_steps, + FusedAdamMathFunctorMP< + float, + float, + float, + scalar_t, + scalar_t, + scalar_t, + 5, + ADAM_MODE::ADAMW, + true>(), + lr_ptr, // unused + lr, + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }), + kBFloat16); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + kHalf, + kBFloat16, + params[0].scalar_type(), + "fused_adamw_kernel_cuda", + [&]() { + multi_tensor_apply_for_fused_optimizer<5>( + tensor_lists, + state_steps, + FusedAdamMathFunctor(), + lr_ptr, // unused + lr, + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }); + } } // The following overload simply has a Tensor lr @@ -87,26 +126,64 @@ void _fused_adamw_amsgrad_cuda_impl_( found_inf.has_value() ? found_inf->data_ptr() : nullptr; const float* lr_ptr = lr.const_data_ptr(); - AT_DISPATCH_FLOATING_TYPES_AND2( - kHalf, - kBFloat16, - params[0].scalar_type(), - "fused_adamw_kernel_cuda", - [&]() { - multi_tensor_apply_for_fused_optimizer<5>( - tensor_lists, - state_steps, - FusedAdamMathFunctor(), - lr_ptr, - 1.0, // unused - beta1, - beta2, - weight_decay, - eps, - maximize, - grad_scale_ptr, - found_inf_ptr); - }); + if (params[0].scalar_type() != exp_avgs[0].scalar_type()) { + validate_mixed_precision_dtypes( + params, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + "Mixed-precision fused AdamW"); + AT_DISPATCH_V2( + exp_avgs[0].scalar_type(), + "fused_adamw_amsgrad_mp_kernel_cuda", + AT_WRAP([&]() { + multi_tensor_apply_for_fused_optimizer<5>( + tensor_lists, + state_steps, + FusedAdamMathFunctorMP< + float, + float, + float, + scalar_t, + scalar_t, + scalar_t, + 5, + ADAM_MODE::ADAMW, + true>(), + lr_ptr, + 1.0, // unused + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }), + kBFloat16); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + kHalf, + kBFloat16, + params[0].scalar_type(), + "fused_adamw_kernel_cuda", + [&]() { + multi_tensor_apply_for_fused_optimizer<5>( + tensor_lists, + state_steps, + FusedAdamMathFunctor(), + lr_ptr, + 1.0, // unused + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }); + } } } // namespace at::native diff --git a/aten/src/ATen/native/cuda/fused_adamw_impl.cu b/aten/src/ATen/native/cuda/fused_adamw_impl.cu index 90318854bec4c..0a65fee7ba84d 100644 --- a/aten/src/ATen/native/cuda/fused_adamw_impl.cu +++ b/aten/src/ATen/native/cuda/fused_adamw_impl.cu @@ -1,6 +1,7 @@ #include #include +#include #include #include #include @@ -31,26 +32,59 @@ void _fused_adamw_cuda_impl_( found_inf.has_value() ? found_inf->data_ptr() : nullptr; const float* lr_ptr = nullptr; - AT_DISPATCH_FLOATING_TYPES_AND2( - kHalf, - kBFloat16, - params[0].scalar_type(), - "fused_adamw_kernel_cuda", - [&]() { - multi_tensor_apply_for_fused_optimizer<4>( - tensor_lists, - state_steps, - FusedAdamMathFunctor(), - lr_ptr, // unused - lr, - beta1, - beta2, - weight_decay, - eps, - maximize, - grad_scale_ptr, - found_inf_ptr); - }); + if (params[0].scalar_type() != exp_avgs[0].scalar_type()) { + validate_mixed_precision_dtypes( + params, grads, exp_avgs, exp_avg_sqs, "Mixed-precision fused AdamW"); + AT_DISPATCH_V2( + exp_avgs[0].scalar_type(), + "fused_adamw_mp_kernel_cuda", + AT_WRAP([&]() { + multi_tensor_apply_for_fused_optimizer<4>( + tensor_lists, + state_steps, + FusedAdamMathFunctorMP< + float, + float, + float, + scalar_t, + scalar_t, + float, + 4, + ADAM_MODE::ADAMW, + false>(), + lr_ptr, // unused + lr, + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }), + kBFloat16); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + kHalf, + kBFloat16, + params[0].scalar_type(), + "fused_adamw_kernel_cuda", + [&]() { + multi_tensor_apply_for_fused_optimizer<4>( + tensor_lists, + state_steps, + FusedAdamMathFunctor(), + lr_ptr, // unused + lr, + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }); + } } // The following overload simply has a Tensor lr @@ -77,26 +111,59 @@ void _fused_adamw_cuda_impl_( found_inf.has_value() ? found_inf->data_ptr() : nullptr; const float* lr_ptr = lr.const_data_ptr(); - AT_DISPATCH_FLOATING_TYPES_AND2( - kHalf, - kBFloat16, - params[0].scalar_type(), - "fused_adamw_kernel_cuda", - [&]() { - multi_tensor_apply_for_fused_optimizer<4>( - tensor_lists, - state_steps, - FusedAdamMathFunctor(), - lr_ptr, - 1.0, // unused - beta1, - beta2, - weight_decay, - eps, - maximize, - grad_scale_ptr, - found_inf_ptr); - }); + if (params[0].scalar_type() != exp_avgs[0].scalar_type()) { + validate_mixed_precision_dtypes( + params, grads, exp_avgs, exp_avg_sqs, "Mixed-precision fused AdamW"); + AT_DISPATCH_V2( + exp_avgs[0].scalar_type(), + "fused_adamw_mp_kernel_cuda", + AT_WRAP([&]() { + multi_tensor_apply_for_fused_optimizer<4>( + tensor_lists, + state_steps, + FusedAdamMathFunctorMP< + float, + float, + float, + scalar_t, + scalar_t, + float, + 4, + ADAM_MODE::ADAMW, + false>(), + lr_ptr, + 1.0, // unused + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }), + kBFloat16); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + kHalf, + kBFloat16, + params[0].scalar_type(), + "fused_adamw_kernel_cuda", + [&]() { + multi_tensor_apply_for_fused_optimizer<4>( + tensor_lists, + state_steps, + FusedAdamMathFunctor(), + lr_ptr, + 1.0, // unused + beta1, + beta2, + weight_decay, + eps, + maximize, + grad_scale_ptr, + found_inf_ptr); + }); + } } } // namespace at::native diff --git a/aten/src/ATen/native/cuda/group_norm_kernel.cu b/aten/src/ATen/native/cuda/group_norm_kernel.cu index 470dd801127d5..23c25554b3a0e 100644 --- a/aten/src/ATen/native/cuda/group_norm_kernel.cu +++ b/aten/src/ATen/native/cuda/group_norm_kernel.cu @@ -26,6 +26,19 @@ namespace { constexpr int kCUDANumThreads = 256; constexpr int kReduceTileSize = 32; +// Reduce across exactly 32 lanes (offsets 16, 8, 4, 2, 1). +// On NVIDIA (warp=32) this is identical to WarpReduceSum. +// On AMD (wavefront=64) this avoids summing across two tile columns +// when the block is (32, 16) and consecutive y-rows share a wavefront. +template +__inline__ __device__ T ReduceSum32(T val) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + val += WARP_SHFL_DOWN(val, offset); + } + return val; +} + template __global__ void RowwiseMomentsCUDAKernel( int64_t N, @@ -52,7 +65,7 @@ __global__ void RowwiseMomentsCUDAKernel( // https://github.com/pytorch/pytorch/pull/13967 __shared__ typename std::aligned_storage< sizeof(WelfordType), - alignof(WelfordType)>::type val_shared[C10_WARP_SIZE]; + alignof(WelfordType)>::type val_shared[C10_WARP_SIZE_UPPER_BOUND]; WelfordType* val_shared_ptr = reinterpret_cast(val_shared); val = cuda_utils::BlockReduce( val, @@ -123,8 +136,8 @@ __global__ void Compute1dBackwardFusedParamsCUDAKernel( sum1 = cuda_utils::WarpReduceSum(sum1); sum2 = cuda_utils::WarpReduceSum(sum2); } else { - __shared__ T_ACC ds_shared[C10_WARP_SIZE]; - __shared__ T_ACC db_shared[C10_WARP_SIZE]; + __shared__ T_ACC ds_shared[C10_WARP_SIZE_UPPER_BOUND]; + __shared__ T_ACC db_shared[C10_WARP_SIZE_UPPER_BOUND]; sum1 = cuda_utils::BlockReduceSum(sum1, ds_shared); sum2 = cuda_utils::BlockReduceSum(sum2, db_shared); } @@ -238,8 +251,11 @@ __global__ void GammaBeta1dBackwardCUDAKernel2( // Do warp reduce for the 1st 16 cols in the tile. T_ACC sum1 = g_shared[threadIdx.x][threadIdx.y]; T_ACC sum2 = b_shared[threadIdx.x][threadIdx.y]; - sum1 = cuda_utils::WarpReduceSum(sum1); - sum2 = cuda_utils::WarpReduceSum(sum2); + // Use ReduceSum32 (not WarpReduceSum) to reduce exactly 32 lanes. + // On AMD wavefront-64, WarpReduceSum would incorrectly sum across two + // tile columns since consecutive y-rows share a wavefront. + sum1 = ReduceSum32(sum1); + sum2 = ReduceSum32(sum2); if (threadIdx.x == 0) { const int64_t c = blockIdx.x * blockDim.x + threadIdx.y; if (c < C) { @@ -255,8 +271,8 @@ __global__ void GammaBeta1dBackwardCUDAKernel2( // Do warp reduce for the 2nd 16 cols in the tile. sum1 = g_shared[threadIdx.x][threadIdx.y + blockDim.y]; sum2 = b_shared[threadIdx.x][threadIdx.y + blockDim.y]; - sum1 = cuda_utils::WarpReduceSum(sum1); - sum2 = cuda_utils::WarpReduceSum(sum2); + sum1 = ReduceSum32(sum1); + sum2 = ReduceSum32(sum2); if (threadIdx.x == 0) { const int64_t c = blockIdx.x * blockDim.x + threadIdx.y + blockDim.y; if (c < C) { @@ -290,8 +306,8 @@ __global__ void ComputeInternalGradientsCUDAKernel( sum1 = cuda_utils::WarpReduceSum(sum1); sum2 = cuda_utils::WarpReduceSum(sum2); } else { - __shared__ T_ACC ds_shared[C10_WARP_SIZE]; - __shared__ T_ACC db_shared[C10_WARP_SIZE]; + __shared__ T_ACC ds_shared[C10_WARP_SIZE_UPPER_BOUND]; + __shared__ T_ACC db_shared[C10_WARP_SIZE_UPPER_BOUND]; sum1 = cuda_utils::BlockReduceSum(sum1, ds_shared); sum2 = cuda_utils::BlockReduceSum(sum2, db_shared); } @@ -333,8 +349,8 @@ __global__ void ComputeBackwardFusedParamsCUDAKernel( sum1 = cuda_utils::WarpReduceSum(sum1); sum2 = cuda_utils::WarpReduceSum(sum2); } else { - __shared__ T_ACC ds_shared[C10_WARP_SIZE]; - __shared__ T_ACC db_shared[C10_WARP_SIZE]; + __shared__ T_ACC ds_shared[C10_WARP_SIZE_UPPER_BOUND]; + __shared__ T_ACC db_shared[C10_WARP_SIZE_UPPER_BOUND]; sum1 = cuda_utils::BlockReduceSum(sum1, ds_shared); sum2 = cuda_utils::BlockReduceSum(sum2, db_shared); } @@ -440,10 +456,11 @@ __global__ void GammaBetaBackwardCUDAKernel2( __syncthreads(); // Do warp reduce for the 1st 16 cols in the tile. + // Use ReduceSum32 for correctness on AMD wavefront-64 (see above). T_ACC sum1 = g_shared[threadIdx.x][threadIdx.y]; T_ACC sum2 = b_shared[threadIdx.x][threadIdx.y]; - sum1 = cuda_utils::WarpReduceSum(sum1); - sum2 = cuda_utils::WarpReduceSum(sum2); + sum1 = ReduceSum32(sum1); + sum2 = ReduceSum32(sum2); if (threadIdx.x == 0) { const int64_t c = blockIdx.x * blockDim.x + threadIdx.y; if (c < C) { @@ -459,8 +476,8 @@ __global__ void GammaBetaBackwardCUDAKernel2( // Do warp reduce for the 2nd 16 cols in the tile. sum1 = g_shared[threadIdx.x][threadIdx.y + blockDim.y]; sum2 = b_shared[threadIdx.x][threadIdx.y + blockDim.y]; - sum1 = cuda_utils::WarpReduceSum(sum1); - sum2 = cuda_utils::WarpReduceSum(sum2); + sum1 = ReduceSum32(sum1); + sum2 = ReduceSum32(sum2); if (threadIdx.x == 0) { const int64_t c = blockIdx.x * blockDim.x + threadIdx.y + blockDim.y; if (c < C) { diff --git a/aten/src/ATen/native/cuda/layer_norm_kernel.cu b/aten/src/ATen/native/cuda/layer_norm_kernel.cu index 04333d7450266..59f5e4a041117 100644 --- a/aten/src/ATen/native/cuda/layer_norm_kernel.cu +++ b/aten/src/ATen/native/cuda/layer_norm_kernel.cu @@ -32,12 +32,6 @@ namespace at::native { namespace { constexpr int kCUDANumThreads = 256; -#ifdef USE_ROCM -// C10_WARP_SIZE is not constexpr for host code. -#define kWarpSize C10_WARP_SIZE -#else -constexpr unsigned int kWarpSize = C10_WARP_SIZE; -#endif constexpr int vec_size = 4; //we could make it dependent on dtype, but that would lead to different results between float and low-p types // aligned vector generates vectorized load/store on CUDA (copy-pasted from MemoryAccess.cuh) @@ -67,7 +61,7 @@ __global__ void RowwiseMomentsCUDAKernel( __shared__ typename std::aligned_storage:: - type val_shared[C10_WARP_SIZE]; + type val_shared[C10_WARP_SIZE_UPPER_BOUND]; WelfordType* val_shared_ptr = reinterpret_cast(val_shared); const int64_t i = blockIdx.x; @@ -662,7 +656,7 @@ blockReduceGammaBetaBackwardsHelper( constexpr int rows_per_thread_y = rows_per_block_y / block_dim_y; int64_t thread_x = blockIdx.x * block_dim_x + threadIdx.x; - int lane_id = (threadIdx.y * blockDim.x + threadIdx.x) & (kWarpSize - 1); + int lane_id = (threadIdx.y * blockDim.x + threadIdx.x) & (C10_WARP_SIZE - 1); int64_t mean_index = M_start + threadIdx.y * rows_per_thread_y; T_ACC warp_mean = 0, warp_rstd = 0; if (lane_id < rows_per_thread_y && mean_index + lane_id < M) { @@ -695,9 +689,9 @@ blockReduceGammaBetaBackwardsHelper( #pragma unroll for (int i = 0; i < rows_per_thread_y; ++i) { - T_ACC rstd_reg = WARP_SHFL(warp_rstd, i, kWarpSize); + T_ACC rstd_reg = WARP_SHFL(warp_rstd, i, C10_WARP_SIZE); if constexpr (!rms_norm){ - T_ACC mean_reg = WARP_SHFL(warp_mean, i, kWarpSize); + T_ACC mean_reg = WARP_SHFL(warp_mean, i, C10_WARP_SIZE); dg_sum += dY_regs[i] * (X_regs[i] - mean_reg) * rstd_reg; db_sum += dY_regs[i]; } else{ @@ -775,7 +769,7 @@ __launch_bounds__(block_dim_x * block_dim_y) T* __restrict__ db) { // This assert is a compile-time check only. constexpr int rows_per_thread_y = rows_per_block_y / block_dim_y; - static_assert(rows_per_thread_y <= kWarpSize); + static_assert(rows_per_thread_y <= C10_WARP_SIZE_LOWER_BOUND); T_ACC dg_sum = 0; T_ACC db_sum = 0; @@ -818,7 +812,7 @@ __launch_bounds__(block_dim_x * block_dim_y) } else { // The caller requested a full reduction so we must reduce across // warps using shared memory and warp shuffles. - static_assert(rows_per_thread_y <= C10_WARP_SIZE); + static_assert(rows_per_thread_y <= C10_WARP_SIZE_LOWER_BOUND); alignas(sizeof(double)) extern __shared__ char s_data1[]; T_ACC* s_data_typed = reinterpret_cast(&s_data1); T_ACC* s_dg; @@ -834,8 +828,8 @@ __launch_bounds__(block_dim_x * block_dim_y) // Load transposed so that a warp holds an entire column // Because block_dim_x != block_dim_y in the general case, we need // some code to handle the general case. - static_assert(block_dim_x * block_dim_y % C10_WARP_SIZE == 0); - constexpr int warps_available_to_reduce = block_dim_x * block_dim_y / C10_WARP_SIZE; + static_assert(block_dim_x * block_dim_y % C10_WARP_SIZE_LOWER_BOUND == 0); + const int warps_available_to_reduce = block_dim_x * block_dim_y / C10_WARP_SIZE; int thread_id = threadIdx.y * block_dim_x + threadIdx.x; int warp_id = thread_id / C10_WARP_SIZE; int lane_id = thread_id & (C10_WARP_SIZE - 1); @@ -848,8 +842,8 @@ __launch_bounds__(block_dim_x * block_dim_y) } #pragma unroll for (unsigned delta = block_dim_y >> 1; delta >= 1; delta >>= 1) { - reg_dg += WARP_SHFL_XOR(reg_dg, delta, kWarpSize); - reg_db += WARP_SHFL_XOR(reg_db, delta, kWarpSize); + reg_dg += WARP_SHFL_XOR(reg_dg, delta, C10_WARP_SIZE); + reg_db += WARP_SHFL_XOR(reg_db, delta, C10_WARP_SIZE); } // Reduce is done. Now write it out to global memory. int64_t out_index = ((int64_t)blockIdx.x) * block_dim_x + i; @@ -955,7 +949,11 @@ void LaunchGammaBetaBackwardCUDAKernel( Tensor* dgamma, Tensor* dbeta, cudaStream_t cuda_stream) { +#ifdef USE_ROCM + constexpr int block_dim_x = 64; +#else constexpr int block_dim_x = 32; +#endif const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount; if (M > 64 * 1024 && N / block_dim_x < sm_count / 2) { // We have a situation where M >> N and N is small. @@ -1010,7 +1008,13 @@ void LaunchGammaBetaBackwardCUDAKernel( } else if (M < 256) { ConfigureAndLaunchGammaBetaBackwardKernel(dY_data, X_data, mean_data, rstd_data, M, N, dgamma, dbeta, cuda_stream); } else { +#ifdef USE_ROCM + // Cap block_dim_y at 16 to keep total threads (64*16=1024) within GPU limits. + // rows_per_thread_y = 256/16 = 16, still within warp size constraint. + ConfigureAndLaunchGammaBetaBackwardKernel(dY_data, X_data, mean_data, rstd_data, M, N, dgamma, dbeta, cuda_stream); +#else ConfigureAndLaunchGammaBetaBackwardKernel(dY_data, X_data, mean_data, rstd_data, M, N, dgamma, dbeta, cuda_stream); +#endif } } } @@ -1617,40 +1621,12 @@ void LayerNormBackwardKernelImplInternal( dbeta_data); C10_CUDA_KERNEL_LAUNCH_CHECK(); } else { - // For small batch size, do colwise reduce directly. - const int part_size = warp_size; - const dim3 threads2(warp_size, 4, 1); - const dim3 blocks2((N + threads2.x - 1) / threads2.x, part_size, 1); - const int nshared2_a = 2 * sizeof(T_ACC) * threads2.y * threads2.y * (threads2.x + 1); - const int nshared2_b = threads2.x * threads2.y * sizeof(T_ACC); - const int nshared2 = nshared2_a > nshared2_b ? nshared2_a : nshared2_b; - - const auto part_grad_dtype = at::toAccumulateType(X.scalar_type(), true); - Tensor part_grad_gamma = at::empty({part_size,N}, gamma.options().dtype(part_grad_dtype)); - Tensor part_grad_beta = at::native::empty_like(part_grad_gamma); - - cuComputePartGradGammaBeta<<>>( - dY_data, - X_data, - M,N, - mean_data, - rstd_data, - part_grad_gamma.template data_ptr(), - part_grad_beta.template data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - const dim3 threads3(warp_size, 8, 1); // Optimization for ROCm - const dim3 blocks3((N + threads3.x - 1) / threads3.x, 1, 1); - const int nshared3 = threads3.x * threads3.y * sizeof(T_ACC); - - cuComputeGradGammaBeta<<>>( - part_grad_gamma.template data_ptr(), - part_grad_beta.template data_ptr(), - part_size, - M,N, - dgamma_data, - dbeta_data); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + // Use the optimized tiled kernel adapted for wavefront-64. + // This replaces the legacy two-pass cuComputePartGradGammaBeta + + // cuComputeGradGammaBeta approach with a single-pass tiled reduction + // that has coalesced memory access and adaptive tile sizing. + LaunchGammaBetaBackwardCUDAKernel( + dY_data, X_data, mean_data, rstd_data, M, N, dgamma, dbeta, cuda_stream); } #else LaunchGammaBetaBackwardCUDAKernel( diff --git a/aten/src/ATen/native/cudnn/BatchNorm.h b/aten/src/ATen/native/cudnn/BatchNorm.h index 3da76c0c16e41..7bd0a18b94de2 100644 --- a/aten/src/ATen/native/cudnn/BatchNorm.h +++ b/aten/src/ATen/native/cudnn/BatchNorm.h @@ -1,6 +1,6 @@ namespace at::native { -TORCH_API size_t +TORCH_CUDA_CU_API size_t _get_cudnn_batch_norm_reserve_space_size(const Tensor& input_t, bool training); } // namespace at::native diff --git a/aten/src/ATen/native/cudnn/RNN.cpp b/aten/src/ATen/native/cudnn/RNN.cpp index 9c12ad5f96b3e..a45fda2b22db4 100644 --- a/aten/src/ATen/native/cudnn/RNN.cpp +++ b/aten/src/ATen/native/cudnn/RNN.cpp @@ -2376,12 +2376,7 @@ struct DropoutState { if (event) { #if !defined(USE_ROCM) // See Note [DropoutState and CUDA graph capture] - cudaStreamCaptureStatus status; - AT_CUDA_CHECK(cudaStreamGetCaptureInfo( - cuda::getCurrentCUDAStream(), &status, &capture_id_last_lock)); - if (status == cudaStreamCaptureStatus::cudaStreamCaptureStatusNone) { - capture_id_last_lock = 0; - } + capture_id_last_lock = at::cuda::currentStreamCaptureId().value_or(0); if (capture_id_last_lock == capture_id_last_unlock) { event->block(cuda::getCurrentCUDAStream()); } @@ -2396,12 +2391,7 @@ struct DropoutState { event->record(); #if !defined(USE_ROCM) // See Note [DropoutState and CUDA graph capture] - cudaStreamCaptureStatus status; - AT_CUDA_CHECK(cudaStreamGetCaptureInfo( - cuda::getCurrentCUDAStream(), &status, &capture_id_last_unlock)); - if (status == cudaStreamCaptureStatus::cudaStreamCaptureStatusNone) { - capture_id_last_unlock = 0; - } + capture_id_last_unlock = at::cuda::currentStreamCaptureId().value_or(0); TORCH_INTERNAL_ASSERT(capture_id_last_unlock == capture_id_last_lock); #endif } diff --git a/aten/src/ATen/native/miopen/LossCTC_miopen.cpp b/aten/src/ATen/native/miopen/LossCTC_miopen.cpp index 21797e7537d59..cf5836b7bb296 100644 --- a/aten/src/ATen/native/miopen/LossCTC_miopen.cpp +++ b/aten/src/ATen/native/miopen/LossCTC_miopen.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #endif // TODO: Remove the condition on AT_ROCM_ENABLED entirely, diff --git a/aten/src/ATen/native/mkldnn/xpu/Conv.cpp b/aten/src/ATen/native/mkldnn/xpu/Conv.cpp index eb8f154dedaf0..d6736750fb000 100644 --- a/aten/src/ATen/native/mkldnn/xpu/Conv.cpp +++ b/aten/src/ATen/native/mkldnn/xpu/Conv.cpp @@ -384,9 +384,10 @@ Tensor _convolution_out( at::MemoryFormat mfmt = is_channels_last_suggested ? get_cl_tag_by_ndim(input.ndimension()) : at::MemoryFormat::Contiguous; - auto bias = bias_r.defined() ? bias_r.contiguous() : bias_r; - input = input.contiguous(mfmt); - weight = weight.contiguous(mfmt); + + auto bias = bias_r.defined() ? make_contiguous_and_aligned(bias_r) : bias_r; + input = make_contiguous_and_aligned(input, mfmt); + weight = make_contiguous_and_aligned(weight, mfmt); check_shape_forward(input, weight, bias, params, true); Tensor output; @@ -591,9 +592,9 @@ std::tuple convolution_backward_overrideable( auto mfmt = is_channels_last_suggested ? get_cl_tag_by_ndim(input_.ndimension()) : at::MemoryFormat::Contiguous; - grad_output_ = grad_output_.contiguous(mfmt); - weight_ = weight_.contiguous(mfmt); - input_ = input_.contiguous(mfmt); + grad_output_ = make_contiguous_and_aligned(grad_output_, mfmt); + weight_ = make_contiguous_and_aligned(weight_, mfmt); + input_ = make_contiguous_and_aligned(input_, mfmt); auto opt = grad_output_.options(); Tensor grad_input; diff --git a/aten/src/ATen/native/mkldnn/xpu/detail/QMatmul.cpp b/aten/src/ATen/native/mkldnn/xpu/detail/QMatmul.cpp index f79dfadd65454..953b3b5a1b2e4 100644 --- a/aten/src/ATen/native/mkldnn/xpu/detail/QMatmul.cpp +++ b/aten/src/ATen/native/mkldnn/xpu/detail/QMatmul.cpp @@ -345,8 +345,14 @@ struct ScaleSpec { int64_t outer_dim, int64_t inner_dim, const std::string& arg_type) const { - if (groups == dnnl::memory::dims{1, 1}) + // TensorWise: mask=0, groups={} -> single scale + if (groups.empty()) { + TORCH_INTERNAL_ASSERT( + mask == 0, + "Empty groups only valid for TensorWise (mask=0), got mask=", + mask); return 1; // tensorwise scaling + } TORCH_CHECK( arg_type == "src" || arg_type == "wei", @@ -399,8 +405,8 @@ inline ScaleSpec make_scale_spec( if (scaling_type == at::blas::ScalingType::TensorWise) { // Scale tensorwise. The same as `--attr-scales=common`. // mask=0 : scale whole tensor - // groups={1, 1}: indicates that there is only one group for scaling - return {0, {1, 1}, dnnl::memory::data_type::f32}; + // groups={}: indicates that there is only one group for scaling + return {0, {}, dnnl::memory::data_type::f32}; } else { // (scaling_type == at::blas::ScalingType::RowWise) // Scale RowWise. The same as `--attr-scales=per_dim_01`. @@ -479,7 +485,7 @@ sycl::event scaled_matmul( // scale_result tensor currently only supports scalar(TensorWise Scaling). bool with_dst_scale = scale_result && scale_result->defined(); if (with_dst_scale) { - op_attr.set_scales(DNNL_ARG_DST, 0, {1}, dnnl::memory::data_type::f32); + op_attr.set_scales(DNNL_ARG_DST, 0, {}, dnnl::memory::data_type::f32); } op_attr.set_scratchpad_mode(dnnl::scratchpad_mode::user); diff --git a/aten/src/ATen/native/mkldnn/xpu/detail/Utils.cpp b/aten/src/ATen/native/mkldnn/xpu/detail/Utils.cpp index a8a6b870ff6b6..ae26798d11640 100644 --- a/aten/src/ATen/native/mkldnn/xpu/detail/Utils.cpp +++ b/aten/src/ATen/native/mkldnn/xpu/detail/Utils.cpp @@ -287,6 +287,27 @@ void undo_broadcast(at::Tensor& tensor) { return; } +bool is_64_bytes_aligned(const at::Tensor& tensor) { + constexpr uintptr_t alignment_byte = 64; + auto data_ptr = reinterpret_cast(tensor.const_data_ptr()); + return (data_ptr % alignment_byte) == 0; +} + +at::Tensor make_contiguous_and_aligned( + const at::Tensor& tensor, + std::optional memory_format) { + at::Tensor out = memory_format.has_value() ? tensor.contiguous(*memory_format) + : tensor.contiguous(); + if (!is_64_bytes_aligned(out)) { + TORCH_WARN( + "Tensor is not 64-byte aligned. Cloning to ensure alignment for oneDNN " + "operations, which incurs a device-to-device copy."); + out = out.clone(); + } + + return out; +} + bool is_onednn_matmul_strides(const at::Tensor& tensor) { // https://oneapi-src.github.io/oneDNN/dev_guide_matmul.html // oneDNN matmul only support 2-dim and 3-dim @@ -300,11 +321,8 @@ bool is_onednn_matmul_strides(const at::Tensor& tensor) { if (tensor.is_contiguous()) return true; - if (tensor.storage_offset() > 0) { - // currently onednn asks 64 byte alignment - constexpr int alignment_byte = 64; - if (reinterpret_cast(tensor.data_ptr()) % alignment_byte > 0) - return false; + if (tensor.storage_offset() > 0 && !is_64_bytes_aligned(tensor)) { + return false; } // the overlapped cases are not supported diff --git a/aten/src/ATen/native/mkldnn/xpu/detail/Utils.h b/aten/src/ATen/native/mkldnn/xpu/detail/Utils.h index 0055fc2f296ad..e841ab8de6583 100644 --- a/aten/src/ATen/native/mkldnn/xpu/detail/Utils.h +++ b/aten/src/ATen/native/mkldnn/xpu/detail/Utils.h @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -46,6 +47,12 @@ void undo_broadcast(at::Tensor& tensor); bool is_onednn_matmul_strides(const at::Tensor& tensor); +bool is_64_bytes_aligned(const at::Tensor& tensor); + +at::Tensor make_contiguous_and_aligned( + const at::Tensor& tensor, + std::optional memory_format = std::nullopt); + bool is_broadcast_from_other_to_self( const at::Tensor& self, const at::Tensor& other); diff --git a/aten/src/ATen/native/mkldnn/xpu/detail/oneDNNContext.h b/aten/src/ATen/native/mkldnn/xpu/detail/oneDNNContext.h index d919e87bfa586..f28aaa131e51e 100644 --- a/aten/src/ATen/native/mkldnn/xpu/detail/oneDNNContext.h +++ b/aten/src/ATen/native/mkldnn/xpu/detail/oneDNNContext.h @@ -13,7 +13,7 @@ namespace at::native::onednn { -TORCH_XPU_API dnnl::memory make_onednn_memory( +dnnl::memory make_onednn_memory( dnnl::memory::desc md, dnnl::engine& engine, void* ptr); @@ -22,7 +22,7 @@ TORCH_XPU_API dnnl::memory make_onednn_memory( bool set_onednn_verbose(int level); // GpuEngineManager singleton -struct TORCH_XPU_API GpuEngineManager { +struct GpuEngineManager { static GpuEngineManager& Instance(); // Singleton dnnl::engine& get_engine( @@ -50,7 +50,7 @@ struct TORCH_XPU_API GpuEngineManager { }; // GpuStreamManager singleton -struct TORCH_XPU_API GpuStreamManager { +struct GpuStreamManager { static GpuStreamManager& Instance(); // Singleton dnnl::stream& get_stream( diff --git a/aten/src/ATen/native/mps/MetalShaderLibrary.h b/aten/src/ATen/native/mps/MetalShaderLibrary.h index c39e1e5928df0..defaadf5b9e58 100644 --- a/aten/src/ATen/native/mps/MetalShaderLibrary.h +++ b/aten/src/ATen/native/mps/MetalShaderLibrary.h @@ -141,7 +141,8 @@ class MetalShaderLibrary { TensorIteratorBase& iter, const std::string& name, const std::optional alpha = std::nullopt, - const std::optional scalar_arg_type = std::nullopt); + const std::optional scalar_arg_type = std::nullopt, + bool supports_vec4 = false); void exec_binary_kernel( TensorIteratorBase& iter, const std::string& name, diff --git a/aten/src/ATen/native/mps/OperationUtils.h b/aten/src/ATen/native/mps/OperationUtils.h index 614fb732b87cd..112f22742ac13 100644 --- a/aten/src/ATen/native/mps/OperationUtils.h +++ b/aten/src/ATen/native/mps/OperationUtils.h @@ -44,6 +44,7 @@ struct MPSScalar { float f; // MPS doesn't support 'double' at::Half h; int64_t i; + uint64_t u; bool b; c10::complex cf; c10::complex ch; @@ -250,7 +251,7 @@ struct MPSKernelCache { __block MPSCachedKernel* cachedKernel = nil; MPSCacheKey hash = std::hash{}(key); dispatch_sync_with_rethrow(serialQueue_, ^() { - if (cache_.count(hash) != 0) { + if (cache_.contains(hash)) { auto& entry = cache_.at(hash); TORCH_INTERNAL_ASSERT_DEBUG_ONLY(key == entry.key_, "Key collision in the MPS cached kernel!\n"); cachedKernel = entry.cachedKernel_; diff --git a/aten/src/ATen/native/mps/OperationUtils.mm b/aten/src/ATen/native/mps/OperationUtils.mm index f449ac4f0e1e6..65067ab796cfe 100644 --- a/aten/src/ATen/native/mps/OperationUtils.mm +++ b/aten/src/ATen/native/mps/OperationUtils.mm @@ -442,8 +442,6 @@ bool isTooLargeForMPSGraph(const Tensor& tensor, bool useMPSStridedAPI) { auto sizes = src.sizes(); auto nStrides = strides.size(); auto nonZeroStrides = src.strides(); - int64_t crtNonZeroStride = 1; - bool hasZeroStrides = false; auto sortedStridesIndices = getSortedStrides(nonZeroStrides); NSMutableArray* sortedStridesShape = [NSMutableArray arrayWithCapacity:nStrides]; @@ -645,6 +643,8 @@ MPSScalar getMPSScalar(const Scalar& scalar, ScalarType type) { case ScalarType::ComplexDouble: return {.size = sizeof(int64_t), .type = type, .value.cf = scalar.to>()}; // Unsigned types + case ScalarType::UInt64: + return {.size = sizeof(uint64_t), .type = type, .value.u = scalar.to()}; case ScalarType::UInt32: return {.size = sizeof(uint32_t), .type = type, .value.i = scalar.to()}; case ScalarType::UInt16: @@ -931,16 +931,21 @@ void executeMPSAllocatorCallback(void* ptr, EventType event) override {} return raw_ptr; } -class BundledShaderLibary : public MetalShaderLibrary { +class BundledShaderLibrary : public MetalShaderLibrary { public: - BundledShaderLibary() : MetalShaderLibrary("") {} + BundledShaderLibrary() : MetalShaderLibrary("") {} protected: id getLibrary() override { if (C10_UNLIKELY(!library)) { auto device = MPSDevice::getInstance()->device(); NSError* error = nil; - library = [device newLibraryWithData:getSectionData("metal_basic") error:&error]; +#ifdef CAN_BUILD_METAL_4 + const auto section_name = is_macos_13_or_newer(MacOSVersion::MACOS_VER_26_0_PLUS) ? "metal_40" : "metal_basic"; +#else + const auto section_name = "metal_basic"; +#endif + library = [device newLibraryWithData:getSectionData(section_name) error:&error]; TORCH_CHECK(library, "Failed to create metal library, error: ", [[error description] UTF8String]); } return library; @@ -976,11 +981,12 @@ static dispatch_data_t getSectionData(const std::string& name) { void MetalShaderLibrary::exec_unary_kernel(TensorIteratorBase& iter, const std::string& name, std::optional alpha, - std::optional scalar_arg_type) { + std::optional scalar_arg_type, + bool supports_vec4) { // Decompose 64-bit tensor into 32-bit ones if (!iter.can_use_32bit_indexing()) { for (auto&& sub_iter : iter.with_32bit_indexing()) { - exec_unary_kernel(sub_iter, name, alpha, scalar_arg_type); + exec_unary_kernel(sub_iter, name, alpha, scalar_arg_type, supports_vec4); } return; } @@ -992,10 +998,12 @@ static dispatch_data_t getSectionData(const std::string& name) { return; } using namespace mps; + bool use_vec4 = + supports_vec4 && iter.is_contiguous() && !alpha.has_value() && at::isFloatingType(iter.common_dtype()); const auto alpha_type = scalar_arg_type.has_value() ? scalar_arg_type.value() : iter.common_dtype(); auto kernel_name = fmt::format("{}_{}_{}_{}{}", name, - iter.is_contiguous() ? "dense" : "strided", + use_vec4 ? "dense_vec4" : (iter.is_contiguous() ? "dense" : "strided"), scalarToMetalTypeString(outputTensor), scalarToMetalTypeString(inputTensor), alpha.has_value() ? fmt::format("_{}", scalarToMetalTypeString(alpha_type)) : ""); @@ -1010,17 +1018,22 @@ static dispatch_data_t getSectionData(const std::string& name) { [computeEncoder setComputePipelineState:cplState]; bind_iter_tensors(computeEncoder, iter); - if (!iter.is_contiguous()) { - mtl_setArgs<2>(computeEncoder, - outputTensor.sizes(), - inputTensor.strides(), - outputTensor.strides(), - inputTensor.ndimension()); - } - if (alpha) { - mtl_setBytes(computeEncoder, getMPSScalar(*alpha, alpha_type), iter.is_contiguous() ? 2 : 6); + if (use_vec4) { + mtl_setBytes(computeEncoder, length, 2); + mtl_dispatch1DJob(computeEncoder, cplState, (length + 3) / 4); + } else { + if (!iter.is_contiguous()) { + mtl_setArgs<2>(computeEncoder, + outputTensor.sizes(), + inputTensor.strides(), + outputTensor.strides(), + inputTensor.ndimension()); + } + if (alpha) { + mtl_setBytes(computeEncoder, getMPSScalar(*alpha, alpha_type), iter.is_contiguous() ? 2 : 6); + } + mtl_dispatch1DJob(computeEncoder, cplState, length); } - mtl_dispatch1DJob(computeEncoder, cplState, length); getMPSProfiler().endProfileKernel(cplState); }); @@ -1101,7 +1114,8 @@ static dispatch_data_t getSectionData(const std::string& name) { } } - const auto alpha_type = scalar_arg_type.has_value() ? scalar_arg_type.value() : iter.common_dtype(); + const auto alpha_type = + scalar_arg_type.has_value() ? scalar_arg_type.value() : (cast_needed ? out.scalar_type() : iter.common_dtype()); const auto alpha_suffix = alpha.has_value() ? fmt::format("_{}", scalarToMetalTypeString(alpha_type)) : ""; std::string kernel_name; @@ -1314,7 +1328,7 @@ static dispatch_data_t getSectionData(const std::string& name) { } MetalShaderLibrary& MetalShaderLibrary::getBundledLibrary() { - static BundledShaderLibary l; + static BundledShaderLibrary l; return l; } diff --git a/aten/src/ATen/native/mps/kernels/ActivationKernel.metal b/aten/src/ATen/native/mps/kernels/ActivationKernel.metal index 7d1f3aa5bacf6..dc892dceea2d0 100644 --- a/aten/src/ATen/native/mps/kernels/ActivationKernel.metal +++ b/aten/src/ATen/native/mps/kernels/ActivationKernel.metal @@ -44,6 +44,27 @@ REGISTER_BINARY_ALPHA_OP(shrink_backward, float, float, float); REGISTER_BINARY_ALPHA_OP(shrink_backward, half, half, half); REGISTER_BINARY_ALPHA_OP(shrink_backward, bfloat, bfloat, bfloat); +struct relu_functor { + template + inline T operator()(const T x) { + return x > T(0) ? x : T(0); + } +}; + +REGISTER_UNARY_OP(relu, float, float); +REGISTER_UNARY_OP(relu, half, half); +REGISTER_UNARY_OP(relu, bfloat, bfloat); +REGISTER_UNARY_OP(relu, long, long); +REGISTER_UNARY_OP(relu, int, int); +REGISTER_UNARY_OP(relu, short, short); +REGISTER_UNARY_OP(relu, char, char); +REGISTER_UNARY_OP(relu, uchar, uchar); +REGISTER_UNARY_OP(relu, bool, bool); + +REGISTER_UNARY_VEC4_OP(relu, float, float); +REGISTER_UNARY_VEC4_OP(relu, half, half); +REGISTER_UNARY_VEC4_OP(relu, bfloat, bfloat); + struct hardsigmoid_functor { template inline T operator()(const T x) { @@ -180,3 +201,37 @@ REGISTER_UNARY_ALPHA_OP(leaky_relu, bfloat, bfloat, bfloat); REGISTER_BINARY_ALPHA_OP(leaky_relu_backward, float, float, float); REGISTER_BINARY_ALPHA_OP(leaky_relu_backward, half, half, half); REGISTER_BINARY_ALPHA_OP(leaky_relu_backward, bfloat, bfloat, bfloat); + +struct silu_functor { + template + inline T operator()(const T x) { + float xf = float(x); + return static_cast(xf / (1.0f + ::metal::precise::exp(-xf))); + } +}; + +REGISTER_UNARY_OP(silu, float, float); +REGISTER_UNARY_OP(silu, half, half); +REGISTER_UNARY_OP(silu, bfloat, bfloat); +REGISTER_UNARY_OP(silu, int, int); +REGISTER_UNARY_OP(silu, short, short); +REGISTER_UNARY_OP(silu, char, char); +REGISTER_UNARY_OP(silu, uchar, uchar); +REGISTER_UNARY_OP(silu, bool, bool); + +REGISTER_UNARY_VEC4_OP(silu, float, float); +REGISTER_UNARY_VEC4_OP(silu, half, half); +REGISTER_UNARY_VEC4_OP(silu, bfloat, bfloat); + +struct silu_backward_functor { + template + inline T operator()(const T grad_output, const T self) { + float sf = float(self); + float sig = 1.0f / (1.0f + ::metal::precise::exp(-sf)); + return static_cast(float(grad_output) * sig * (1.0f + sf - sf * sig)); + } +}; + +REGISTER_BINARY_OP(silu_backward, float, float); +REGISTER_BINARY_OP(silu_backward, half, half); +REGISTER_BINARY_OP(silu_backward, bfloat, bfloat); diff --git a/aten/src/ATen/native/mps/kernels/BinaryKernel.metal b/aten/src/ATen/native/mps/kernels/BinaryKernel.metal index eb7deb62a9042..1def63fff16ae 100644 --- a/aten/src/ATen/native/mps/kernels/BinaryKernel.metal +++ b/aten/src/ATen/native/mps/kernels/BinaryKernel.metal @@ -560,3 +560,176 @@ REGISTER_BINARY_ALPHA_OP(sub_alpha, float2, float2, float2); REGISTER_BINARY_ALPHA_OP(sub_alpha, half2, half2, half2); REGISTER_BINARY_ALPHA_OP(lerp_alpha, float2, float2, float2); REGISTER_BINARY_ALPHA_OP(lerp_alpha, half2, half2, half2); + +// lerp with tensor weight: lerp(s, e, w) = fma(w, e - s, s) +template +inline T lerp_op(T s, T e, T w) { + return fma(w, e - s, s); +} + +inline bfloat lerp_op(bfloat s, bfloat e, bfloat w) { + return static_cast(fma(float(w), float(e) - float(s), float(s))); +} + +inline long lerp_op(long s, long e, long w) { + return s + w * (e - s); +} + +inline float2 lerp_op(float2 s, float2 e, float2 w) { + return s + mul(w, e - s); +} + +template +kernel void lerp_tensor_dense( + device T* out [[buffer(0)]], + device const T* self [[buffer(1)]], + device const T* end [[buffer(2)]], + device const T* weight [[buffer(3)]], + uint tid [[thread_position_in_grid]]) { + out[tid] = lerp_op(self[tid], end[tid], weight[tid]); +} + +// Scalar weight broadcast: self/end/out contiguous, weight is a single element +template +kernel void lerp_tensor_scalar_weight( + device T* out [[buffer(0)]], + device const T* self [[buffer(1)]], + device const T* end [[buffer(2)]], + device const T& weight [[buffer(3)]], + uint tid [[thread_position_in_grid]]) { + out[tid] = lerp_op(self[tid], end[tid], weight); +} + +// 2D strided: coordinates from 2D dispatch, no integer division +template +kernel void lerp_tensor_strided_2d( + device void* out_ptr [[buffer(0)]], + constant void* self_ptr [[buffer(1)]], + constant void* end_ptr [[buffer(2)]], + constant void* weight_ptr [[buffer(3)]], + constant long* out_strides [[buffer(4)]], + constant long* self_strides [[buffer(5)]], + constant long* end_strides [[buffer(6)]], + constant long* weight_strides [[buffer(7)]], + uint2 tid [[thread_position_in_grid]]) { + int out_off = + int(tid.x) * int(out_strides[0]) + int(tid.y) * int(out_strides[1]); + int self_off = + int(tid.x) * int(self_strides[0]) + int(tid.y) * int(self_strides[1]); + int end_off = + int(tid.x) * int(end_strides[0]) + int(tid.y) * int(end_strides[1]); + int wt_off = + int(tid.x) * int(weight_strides[0]) + int(tid.y) * int(weight_strides[1]); + ref_at_offs(out_ptr, long(out_off)) = lerp_op( + val_at_offs(self_ptr, long(self_off)), + val_at_offs(end_ptr, long(end_off)), + val_at_offs(weight_ptr, long(wt_off))); +} + +// 3D strided: coordinates from 3D dispatch, no integer division +template +kernel void lerp_tensor_strided_3d( + device void* out_ptr [[buffer(0)]], + constant void* self_ptr [[buffer(1)]], + constant void* end_ptr [[buffer(2)]], + constant void* weight_ptr [[buffer(3)]], + constant long* out_strides [[buffer(4)]], + constant long* self_strides [[buffer(5)]], + constant long* end_strides [[buffer(6)]], + constant long* weight_strides [[buffer(7)]], + uint3 tid [[thread_position_in_grid]]) { + int out_off = int(tid.x) * int(out_strides[0]) + + int(tid.y) * int(out_strides[1]) + int(tid.z) * int(out_strides[2]); + int self_off = int(tid.x) * int(self_strides[0]) + + int(tid.y) * int(self_strides[1]) + int(tid.z) * int(self_strides[2]); + int end_off = int(tid.x) * int(end_strides[0]) + + int(tid.y) * int(end_strides[1]) + int(tid.z) * int(end_strides[2]); + int wt_off = int(tid.x) * int(weight_strides[0]) + + int(tid.y) * int(weight_strides[1]) + int(tid.z) * int(weight_strides[2]); + ref_at_offs(out_ptr, long(out_off)) = lerp_op( + val_at_offs(self_ptr, long(self_off)), + val_at_offs(end_ptr, long(end_off)), + val_at_offs(weight_ptr, long(wt_off))); +} + +template +kernel void lerp_tensor_strided( + device void* out_ptr [[buffer(0)]], + constant void* self_ptr [[buffer(1)]], + constant void* end_ptr [[buffer(2)]], + constant void* weight_ptr [[buffer(3)]], + constant long* sizes [[buffer(4)]], + constant long* out_strides [[buffer(5)]], + constant long* self_strides [[buffer(6)]], + constant long* end_strides [[buffer(7)]], + constant long* weight_strides [[buffer(8)]], + constant uint& ndim [[buffer(9)]], + uint tid [[thread_position_in_grid]]) { + int pos[max_ndim]; + pos_from_thread_index(int(tid), pos, sizes, ndim); + auto self_off = offset_from_coord(pos, self_strides, ndim); + auto end_off = offset_from_coord(pos, end_strides, ndim); + auto weight_off = offset_from_coord(pos, weight_strides, ndim); + auto out_off = offset_from_coord(pos, out_strides, ndim); + ref_at_offs(out_ptr, out_off) = lerp_op( + val_at_offs(self_ptr, self_off), + val_at_offs(end_ptr, end_off), + val_at_offs(weight_ptr, weight_off)); +} + +#define INSTANTIATE_LERP(DTYPE) \ + template [[host_name("lerp_tensor_dense_" #DTYPE)]] kernel void \ + lerp_tensor_dense( \ + device DTYPE * out [[buffer(0)]], \ + device const DTYPE* self [[buffer(1)]], \ + device const DTYPE* end [[buffer(2)]], \ + device const DTYPE* weight [[buffer(3)]], \ + uint tid [[thread_position_in_grid]]); \ + template [[host_name("lerp_tensor_scalar_weight_" #DTYPE)]] kernel void \ + lerp_tensor_scalar_weight( \ + device DTYPE * out [[buffer(0)]], \ + device const DTYPE* self [[buffer(1)]], \ + device const DTYPE* end [[buffer(2)]], \ + device const DTYPE& weight [[buffer(3)]], \ + uint tid [[thread_position_in_grid]]); \ + template [[host_name("lerp_tensor_strided_2d_" #DTYPE)]] kernel void \ + lerp_tensor_strided_2d( \ + device void* out_ptr [[buffer(0)]], \ + constant void* self_ptr [[buffer(1)]], \ + constant void* end_ptr [[buffer(2)]], \ + constant void* weight_ptr [[buffer(3)]], \ + constant long* out_strides [[buffer(4)]], \ + constant long* self_strides [[buffer(5)]], \ + constant long* end_strides [[buffer(6)]], \ + constant long* weight_strides [[buffer(7)]], \ + uint2 tid [[thread_position_in_grid]]); \ + template [[host_name("lerp_tensor_strided_3d_" #DTYPE)]] kernel void \ + lerp_tensor_strided_3d( \ + device void* out_ptr [[buffer(0)]], \ + constant void* self_ptr [[buffer(1)]], \ + constant void* end_ptr [[buffer(2)]], \ + constant void* weight_ptr [[buffer(3)]], \ + constant long* out_strides [[buffer(4)]], \ + constant long* self_strides [[buffer(5)]], \ + constant long* end_strides [[buffer(6)]], \ + constant long* weight_strides [[buffer(7)]], \ + uint3 tid [[thread_position_in_grid]]); \ + template [[host_name("lerp_tensor_strided_" #DTYPE)]] kernel void \ + lerp_tensor_strided( \ + device void* out_ptr [[buffer(0)]], \ + constant void* self_ptr [[buffer(1)]], \ + constant void* end_ptr [[buffer(2)]], \ + constant void* weight_ptr [[buffer(3)]], \ + constant long* sizes [[buffer(4)]], \ + constant long* out_strides [[buffer(5)]], \ + constant long* self_strides [[buffer(6)]], \ + constant long* end_strides [[buffer(7)]], \ + constant long* weight_strides [[buffer(8)]], \ + constant uint& ndim [[buffer(9)]], \ + uint tid [[thread_position_in_grid]]); + +INSTANTIATE_LERP(float); +INSTANTIATE_LERP(half); +INSTANTIATE_LERP(bfloat); +INSTANTIATE_LERP(float2); +INSTANTIATE_LERP(long); diff --git a/aten/src/ATen/native/mps/kernels/ConstantKernel.metal b/aten/src/ATen/native/mps/kernels/ConstantKernel.metal new file mode 100644 index 0000000000000..1b4c1277b6448 --- /dev/null +++ b/aten/src/ATen/native/mps/kernels/ConstantKernel.metal @@ -0,0 +1,102 @@ +#include +#include + +using namespace metal; + +template +kernel void fill_scalar_dense( + device T* out [[buffer(0)]], + constant T& fill_val [[buffer(1)]], + constant uint& numel [[buffer(2)]], + uint index [[thread_position_in_grid]]) { + if (index < numel) + out[index] = fill_val; +} + +// Single-byte types: each thread fills 4 elements via vec writes. +template +kernel void fill_scalar_dense_vec4( + device T* out [[buffer(0)]], + constant T& fill_val [[buffer(1)]], + constant uint& numel [[buffer(2)]], + uint index [[thread_position_in_grid]]) { + uint base = index * 4; + if (base + 4 <= numel) { + *(device vec*)(out + base) = vec(fill_val); + } else { + for (uint i = base; i < numel; i++) + out[i] = fill_val; + } +} + +#define REGISTER_FILL_OP(T) \ + template [[host_name("fill_scalar_dense_" #T)]] kernel void \ + fill_scalar_dense(device T*, constant T&, constant uint&, uint) + +REGISTER_FILL_OP(float); +REGISTER_FILL_OP(half); +REGISTER_FILL_OP(bfloat); +REGISTER_FILL_OP(long); +REGISTER_FILL_OP(ulong); +REGISTER_FILL_OP(int); +REGISTER_FILL_OP(uint); +REGISTER_FILL_OP(short); +REGISTER_FILL_OP(ushort); +REGISTER_FILL_OP(float2); +REGISTER_FILL_OP(half2); + +// Byte types use the vec4 variant under the same naming convention. +#define REGISTER_FILL_BYTE_OP(T) \ + template [[host_name("fill_scalar_dense_" #T)]] kernel void \ + fill_scalar_dense_vec4(device T*, constant T&, constant uint&, uint) + +REGISTER_FILL_BYTE_OP(char); +REGISTER_FILL_BYTE_OP(uchar); +REGISTER_FILL_BYTE_OP(bool); + +// 2D dispatch: tid.y = dim-0 index (no division), tid.x = linear index for +// dims 1..ndim-1. For an N-dim tensor this requires N-1 divisions instead of N, +// and consecutive threads in x access consecutive addresses in the innermost +// dimension (coalesced writes). +// Strides from TensorIterator are in bytes, use ref_at_offs(out+offs) +template +kernel void fill_scalar_strided( + device void* out [[buffer(0)]], + constant T& fill_val [[buffer(1)]], + constant long* sizes [[buffer(2)]], + constant long* strides [[buffer(3)]], + constant uint& ndim [[buffer(4)]], + uint2 tid [[thread_position_in_grid]]) { + long offset = long(tid.y) * strides[0]; + uint inner = tid.x; + for (uint i = 1; i < ndim; i++) { + offset += long(inner % uint(sizes[i])) * strides[i]; + inner /= uint(sizes[i]); + } + c10::metal::ref_at_offs(out, offset) = fill_val; +} + +#define REGISTER_FILL_STRIDED_OP(T) \ + template [[host_name("fill_scalar_strided_" #T)]] kernel void \ + fill_scalar_strided( \ + device void*, \ + constant T&, \ + constant long*, \ + constant long*, \ + constant uint&, \ + uint2) + +REGISTER_FILL_STRIDED_OP(float); +REGISTER_FILL_STRIDED_OP(half); +REGISTER_FILL_STRIDED_OP(bfloat); +REGISTER_FILL_STRIDED_OP(long); +REGISTER_FILL_STRIDED_OP(ulong); +REGISTER_FILL_STRIDED_OP(int); +REGISTER_FILL_STRIDED_OP(uint); +REGISTER_FILL_STRIDED_OP(short); +REGISTER_FILL_STRIDED_OP(ushort); +REGISTER_FILL_STRIDED_OP(char); +REGISTER_FILL_STRIDED_OP(uchar); +REGISTER_FILL_STRIDED_OP(bool); +REGISTER_FILL_STRIDED_OP(float2); +REGISTER_FILL_STRIDED_OP(half2); diff --git a/aten/src/ATen/native/mps/kernels/Distributions.metal b/aten/src/ATen/native/mps/kernels/Distributions.metal index 6c9a2722f378d..2c8a9404ab533 100644 --- a/aten/src/ATen/native/mps/kernels/Distributions.metal +++ b/aten/src/ATen/native/mps/kernels/Distributions.metal @@ -1,4 +1,5 @@ #include +#include #include using namespace metal; @@ -49,6 +50,26 @@ kernel void geometric( output[tid] = static_cast(result); } +template +kernel void exponential( + device T* output [[buffer(0)]], + constant float2& params [[buffer(1)]], + constant long2& seed_base_offset [[buffer(2)]], + constant uint& numel [[buffer(3)]], + uint tid [[thread_position_in_grid]]) { + uint base = tid * 4; + uint4 raw = + c10::metal::philox4::rand(seed_base_offset.x, seed_base_offset.y + tid); + float lambda = params.x; + uint count = min(4u, numel - base); + for (uint i = 0; i < count; ++i) { + float u = clamp( + c10::metal::detail::uint32_to_uniform_float(raw[i]), eps, 1.0f - eps); + output[base + i] = + static_cast(-::metal::precise::log(1.0f - u) / lambda); + } +} + #define REGISTER_OP(NAME, DTYPE) \ template [[host_name(#NAME "_" #DTYPE)]] kernel void NAME( \ device DTYPE*, constant float2&, constant long2&, uint) @@ -69,3 +90,190 @@ REGISTER_OP(geometric, long); REGISTER_OP(geometric, short); REGISTER_OP(geometric, char); REGISTER_OP(geometric, uchar); + +#define REGISTER_EXPONENTIAL(DTYPE) \ + template [[host_name("exponential_" #DTYPE)]] kernel void \ + exponential( \ + device DTYPE*, constant float2&, constant long2&, constant uint&, uint) + +REGISTER_EXPONENTIAL(float); +REGISTER_EXPONENTIAL(half); +REGISTER_EXPONENTIAL(bfloat); + +// Marsaglia & Tsang (2000) acceptance-rejection method for Gamma distribution. +// Adapted from aten/src/ATen/native/Distributions.h sample_gamma(), +// which originates from NumPy's random module (Copyright 2005 Robert Kern). +// Each thread uses a per-thread RNG offset stride to allow variable-length +// rejection loops without colliding with other threads' random streams. +constant constexpr int GAMMA_RANDOMS_STRIDE = 32; + +template +kernel void standard_gamma( + device T* output [[buffer(0)]], + device const T* alpha_in [[buffer(1)]], + constant long2& seed_base_offset [[buffer(2)]], + uint tid [[thread_position_in_grid]]) { + float alpha = static_cast(alpha_in[tid]); + float scale = 1.0f; + long base = + seed_base_offset.y + static_cast(tid) * GAMMA_RANDOMS_STRIDE; + long seed = seed_base_offset.x; + int rng_idx = 0; + + // Boost alpha < 1 for higher acceptance probability + if (alpha < 1.0f) { + if (alpha == 0.0f) { + output[tid] = static_cast(0.0f); + return; + } + float u = c10::metal::rand(seed, base + rng_idx++); + scale = ::metal::precise::pow(1.0f - u, 1.0f / alpha); + alpha += 1.0f; + } + + // Marsaglia & Tsang acceptance-rejection + float d = alpha - 1.0f / 3.0f; + float c = 1.0f / ::metal::precise::sqrt(9.0f * d); + for (;;) { + float x, y; + do { + x = c10::metal::randn(seed, base + rng_idx++); + y = 1.0f + c * x; + } while (y <= 0.0f); + float v = y * y * y; + float u = 1.0f - c10::metal::rand(seed, base + rng_idx++); + float xx = x * x; + if (u < 1.0f - 0.0331f * xx * xx) { + float result = scale * d * v; + output[tid] = static_cast(max(result, FLT_MIN)); + return; + } + if (::metal::precise::log(u) < + 0.5f * xx + d * (1.0f - v + ::metal::precise::log(v))) { + float result = scale * d * v; + output[tid] = static_cast(max(result, FLT_MIN)); + return; + } + } +} + +#define REGISTER_GAMMA(DTYPE) \ + template [[host_name("standard_gamma_" #DTYPE)]] \ + kernel void standard_gamma( \ + device DTYPE*, device const DTYPE*, constant long2&, uint) + +REGISTER_GAMMA(float); +REGISTER_GAMMA(half); +REGISTER_GAMMA(bfloat); + +// Reparameterized gradient for Gamma distribution. +// Computes -(d/dalpha cdf(x;alpha)) / pdf(x;alpha). +// Adapted from aten/src/ATen/native/Distributions.h standard_gamma_grad_one(). + +constant constexpr float GAMMA_GRAD_COEF_UV[3][8] = { + {0.16009398f, + -0.094634809f, + 0.025146376f, + -0.0030648343f, + 1.0f, + 0.32668115f, + 0.10406089f, + 0.0014179084f}, + {0.53487893f, + 0.1298071f, + 0.065735949f, + -0.0015649758f, + 0.16639465f, + 0.020070113f, + -0.0035938915f, + -0.00058392623f}, + {0.040121004f, + -0.0065914022f, + -0.0026286047f, + -0.0013441777f, + 0.017050642f, + -0.0021309326f, + 0.00085092367f, + -1.5247877e-07f}, +}; + +template +kernel void standard_gamma_grad( + device T* output [[buffer(0)]], + device const T* self_data [[buffer(1)]], + device const T* output_data [[buffer(2)]], + uint tid [[thread_position_in_grid]]) { + float alpha = static_cast(self_data[tid]); + float x = static_cast(output_data[tid]); + + // Region 1: Small x - Taylor series expansion + if (x < 0.8f) { + float numer = 1.0f; + float denom = alpha; + float series1 = numer / denom; + float series2 = numer / (denom * denom); + for (int i = 1; i <= 5; ++i) { + numer *= -x / static_cast(i); + denom += 1.0f; + series1 += numer / denom; + series2 += numer / (denom * denom); + } + float pow_x_alpha = ::metal::precise::pow(x, alpha); + float gamma_pdf = + ::metal::precise::pow(x, alpha - 1.0f) * ::metal::precise::exp(-x); + float gamma_cdf = pow_x_alpha * series1; + float gamma_cdf_alpha = + (::metal::precise::log(x) - c10::metal::digamma(alpha)) * gamma_cdf - + pow_x_alpha * series2; + float result = -gamma_cdf_alpha / gamma_pdf; + output[tid] = static_cast(isnan(result) ? 0.0f : result); + return; + } + + // Region 2: Large alpha - Rice saddle point expansion + if (alpha > 8.0f) { + if (0.9f * alpha <= x && x <= 1.1f * alpha) { + float numer_1 = 1.0f + 24.0f * alpha * (1.0f + 12.0f * alpha); + float numer_2 = 1440.0f * (alpha * alpha) + + 6.0f * x * (53.0f - 120.0f * x) - 65.0f * x * x / alpha + + alpha * (107.0f + 3600.0f * x); + float denom = 1244160.0f * (alpha * alpha) * (alpha * alpha); + output[tid] = static_cast(numer_1 * numer_2 / denom); + return; + } + float denom = ::metal::precise::sqrt(8.0f * alpha); + float term2 = denom / (alpha - x); + float term3 = ::metal::precise::pow( + x - alpha - alpha * ::metal::precise::log(x / alpha), -1.5f); + float term23 = (x < alpha) ? term2 - term3 : term2 + term3; + float term1 = ::metal::precise::log(x / alpha) * term23 - + ::metal::precise::sqrt(2.0f / alpha) * (alpha + x) / + ((alpha - x) * (alpha - x)); + float stirling = + 1.0f + 1.0f / (12.0f * alpha) * (1.0f + 1.0f / (24.0f * alpha)); + float numer = x * term1; + output[tid] = static_cast(-stirling * numer / denom); + return; + } + + // Region 3: Moderate alpha - bivariate rational approximation + float u = ::metal::precise::log(x / alpha); + float v = ::metal::precise::log(alpha); + float coef_v[8]; + for (int i = 0; i < 8; ++i) { + coef_v[i] = GAMMA_GRAD_COEF_UV[0][i] + + u * (GAMMA_GRAD_COEF_UV[1][i] + u * GAMMA_GRAD_COEF_UV[2][i]); + } + float p = coef_v[0] + v * (coef_v[1] + v * (coef_v[2] + v * coef_v[3])); + float q = coef_v[4] + v * (coef_v[5] + v * (coef_v[6] + v * coef_v[7])); + output[tid] = static_cast(::metal::precise::exp(p / q)); +} + +#define REGISTER_GAMMA_GRAD(DTYPE) \ + template [[host_name("standard_gamma_grad_" #DTYPE)]] \ + kernel void standard_gamma_grad( \ + device DTYPE*, device const DTYPE*, device const DTYPE*, uint) + +REGISTER_GAMMA_GRAD(float); +REGISTER_GAMMA_GRAD(half); +REGISTER_GAMMA_GRAD(bfloat); diff --git a/aten/src/ATen/native/mps/kernels/Eye.metal b/aten/src/ATen/native/mps/kernels/Eye.metal new file mode 100644 index 0000000000000..ea389866241c4 --- /dev/null +++ b/aten/src/ATen/native/mps/kernels/Eye.metal @@ -0,0 +1,59 @@ +#include +using namespace metal; + +// For complex types (float2/half2), "1" is (1, 0), not (1, 1) +template +constexpr T eye_one() { + return static_cast(1); +} +template <> +constexpr float2 eye_one() { + return float2(1.0, 0.0); +} +template <> +constexpr half2 eye_one() { + return half2(half(1.0), half(0.0)); +} + +// Single-pass: writes both 0s and 1s in one dispatch (better for small tensors) +template +kernel void eye( + device T* output [[buffer(0)]], + constant long& stride0 [[buffer(1)]], + constant long& stride1 [[buffer(2)]], + uint2 pos [[thread_position_in_grid]]) { + output[pos.y * stride0 + pos.x * stride1] = + (pos.x == pos.y) ? eye_one() : static_cast(0); +} + +// Diagonal-only: writes 1s to pre-zeroed tensor (better for large tensors) +template +kernel void eye_diag( + device T* output [[buffer(0)]], + constant long& diag_stride [[buffer(1)]], + uint index [[thread_position_in_grid]]) { + output[index * diag_stride] = eye_one(); +} + +#define REGISTER_EYE_OP(DTYPE) \ + template [[host_name("eye_" #DTYPE)]] kernel void eye( \ + device DTYPE * output [[buffer(0)]], \ + constant long& stride0 [[buffer(1)]], \ + constant long& stride1 [[buffer(2)]], \ + uint2 pos [[thread_position_in_grid]]); \ + template [[host_name("eye_diag_" #DTYPE)]] kernel void eye_diag( \ + device DTYPE * output [[buffer(0)]], \ + constant long& diag_stride [[buffer(1)]], \ + uint index [[thread_position_in_grid]]); + +REGISTER_EYE_OP(float); +REGISTER_EYE_OP(half); +REGISTER_EYE_OP(bfloat); +REGISTER_EYE_OP(float2); +REGISTER_EYE_OP(half2); +REGISTER_EYE_OP(long); +REGISTER_EYE_OP(int); +REGISTER_EYE_OP(short); +REGISTER_EYE_OP(char); +REGISTER_EYE_OP(uchar); +REGISTER_EYE_OP(bool); diff --git a/aten/src/ATen/native/mps/kernels/GridSampler.metal b/aten/src/ATen/native/mps/kernels/GridSampler.metal index 0d8e6aab87e9e..8314f63bdbce9 100644 --- a/aten/src/ATen/native/mps/kernels/GridSampler.metal +++ b/aten/src/ATen/native/mps/kernels/GridSampler.metal @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -148,36 +149,6 @@ struct PadReflection { } }; -// Cubic convolution helper 1: for |x| < 1 -template -static T cubic_convolution1(T x, T A) { - return ((A + 2) * x - (A + 3)) * x * x + 1; -} - -// Cubic convolution helper 2: for 1 <= |x| < 2 -template -static T cubic_convolution2(T x, T A) { - return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A; -} - -// Get cubic upsampling coefficients (Catmull-Rom spline with A=-0.75) -template -static void get_cubic_coefficients(T coeffs[4], T t) { - T A = static_cast(-0.75); - coeffs[0] = cubic_convolution2(t + 1, A); - coeffs[1] = cubic_convolution1(t, A); - coeffs[2] = cubic_convolution1(1 - t, A); - coeffs[3] = cubic_convolution2(2 - t, A); -} - -// 1D cubic interpolation -template -static T cubic_interp1d(T x0, T x1, T x2, T x3, T t) { - T coeffs[4]; - get_cubic_coefficients(coeffs, t); - return x0 * coeffs[0] + x1 * coeffs[1] + x2 * coeffs[2] + x3 * coeffs[3]; -} - // 2D Bilinear interpolation template static T interpolate_bilinear_2d( diff --git a/aten/src/ATen/native/mps/kernels/Indexing.metal b/aten/src/ATen/native/mps/kernels/Indexing.metal index 463c1812d0302..dcc0c4ff3cb82 100644 --- a/aten/src/ATen/native/mps/kernels/Indexing.metal +++ b/aten/src/ATen/native/mps/kernels/Indexing.metal @@ -760,14 +760,6 @@ INSTANTIATE_INDEX_FILL(uchar); INSTANTIATE_INDEX_FILL(float2); INSTANTIATE_INDEX_FILL(half2); -[[host_name("index_fill_set_mask")]] -kernel void index_fill_set_mask( - device bool*, - constant long*, - constant long&, - constant long&, - uint); - #define INSTANTIATE_INDEX_FILL_FROM_MASK(T) \ template [[host_name("index_fill_dense_from_mask_" #T)]] \ kernel void index_fill_dense_from_mask( \ @@ -799,3 +791,237 @@ INSTANTIATE_INDEX_FILL_FROM_MASK(uchar) INSTANTIATE_INDEX_FILL_FROM_MASK(bool) INSTANTIATE_INDEX_FILL_FROM_MASK(float2) INSTANTIATE_INDEX_FILL_FROM_MASK(half2) + +// Nonzero kernel implementation using prefix-sum + scatter approach. +// +// Step 1 (count_nonzero_prefix_sum): Each threadgroup computes an exclusive +// prefix sum of the nonzero flags over its chunk. Per-threadgroup totals are +// written to block_sums. +// +// Step 2 (prefix_sum_blocks): A single threadgroup computes the exclusive +// prefix sum of block_sums → block_offsets and writes the total nonzero count +// to a 1-element buffer. The host reads back only that single int, then +// allocates the output tensor. +// +// Step 3 (scatter_nonzero_indices): Each thread with a nonzero element writes +// its multi-dimensional indices into the output at the position determined by +// block_offsets[tgid] + prefix[tid]. + +template , bool> = true> +inline bool is_nonzero(T val) { + return val != T(0); +} + +template , bool> = true> +inline bool is_nonzero(T val) { + return val.x != 0 || val.y != 0; +} + +template +[[max_total_threads_per_threadgroup(1024)]] +kernel void count_nonzero_prefix_sum( + const device T* input [[buffer(0)]], + device int* prefix [[buffer(1)]], + device int* block_sums [[buffer(2)]], + uint tid [[thread_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint tgsize [[threads_per_threadgroup]], + uint tgid [[threadgroup_position_in_grid]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]]) { + uint num_simds = (tgsize + simdgroup_size - 1) / simdgroup_size; + + int flag = is_nonzero(input[tid]) ? 1 : 0; + + // Inclusive prefix sum within SIMD group using shuffle + int val = flag; + for (uint offset = 1; offset < simdgroup_size; offset <<= 1) { + int other = simd_shuffle_and_fill_up(val, 0, static_cast(offset)); + val += other; + } + + // The last lane in each simd group writes its total. + // For full groups this is lane 31; for the last (partial) group we compute + // which lane is actually last. + threadgroup int simdgroup_totals[32]; + bool is_last_lane_in_simd; + if (simd_group_id < num_simds - 1) { + is_last_lane_in_simd = (simd_lane_id == simdgroup_size - 1); + } else { + uint lanes_in_last = tgsize - simd_group_id * simdgroup_size; + is_last_lane_in_simd = (simd_lane_id == lanes_in_last - 1); + } + if (is_last_lane_in_simd) { + simdgroup_totals[simd_group_id] = val; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // First simd group computes exclusive prefix sum of simd group totals + threadgroup int simdgroup_offsets[32]; + if (simd_group_id == 0) { + int sg_val = + (simd_lane_id < num_simds) ? simdgroup_totals[simd_lane_id] : 0; + for (uint offset = 1; offset < simdgroup_size; offset <<= 1) { + int other = + simd_shuffle_and_fill_up(sg_val, 0, static_cast(offset)); + sg_val += other; + } + int exclusive = simd_shuffle_and_fill_up(sg_val, 0, static_cast(1)); + simdgroup_offsets[simd_lane_id] = exclusive; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + int exclusive_val = val - flag + simdgroup_offsets[simd_group_id]; + + prefix[tid] = exclusive_val; + + if (lid == tgsize - 1) { + block_sums[tgid] = + simdgroup_offsets[num_simds - 1] + simdgroup_totals[num_simds - 1]; + } +} + +// Step 2: exclusive prefix sum of block_sums → block_offsets, and write +// total nonzero count to a 1-element buffer. Runs in a single threadgroup. +// Each thread handles ceil(num_blocks / tgsize) consecutive blocks via a +// serial loop, then the per-thread totals are scanned in parallel. +[[max_total_threads_per_threadgroup(1024)]] +kernel void prefix_sum_blocks( + const device int* block_sums [[buffer(0)]], + device int* block_offsets [[buffer(1)]], + device int* total_nonzero [[buffer(2)]], + constant uint& num_blocks [[buffer(3)]], + uint lid [[thread_position_in_threadgroup]], + uint tgsize [[threads_per_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]], + uint simd_group_id [[simdgroup_index_in_threadgroup]]) { + uint num_simds = (tgsize + simdgroup_size - 1) / simdgroup_size; + + // Each thread handles a contiguous chunk of blocks + uint chunk_size = (num_blocks + tgsize - 1) / tgsize; + uint start = lid * chunk_size; + uint end = min(start + chunk_size, num_blocks); + + // Serial sum over this thread's chunk + int chunk_total = 0; + for (uint i = start; i < end; i++) { + chunk_total += block_sums[i]; + } + + // Parallel inclusive prefix sum of chunk_totals across threads + int val = chunk_total; + for (uint offset = 1; offset < simdgroup_size; offset <<= 1) { + int other = simd_shuffle_and_fill_up(val, 0, static_cast(offset)); + val += other; + } + + threadgroup int simdgroup_totals[32]; + bool is_last_lane_in_simd; + if (simd_group_id < num_simds - 1) { + is_last_lane_in_simd = (simd_lane_id == simdgroup_size - 1); + } else { + uint lanes_in_last = tgsize - simd_group_id * simdgroup_size; + is_last_lane_in_simd = (simd_lane_id == lanes_in_last - 1); + } + if (is_last_lane_in_simd) { + simdgroup_totals[simd_group_id] = val; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup int simdgroup_offsets[32]; + if (simd_group_id == 0) { + int sg_val = + (simd_lane_id < num_simds) ? simdgroup_totals[simd_lane_id] : 0; + for (uint offset = 1; offset < simdgroup_size; offset <<= 1) { + int other = + simd_shuffle_and_fill_up(sg_val, 0, static_cast(offset)); + sg_val += other; + } + int exclusive = simd_shuffle_and_fill_up(sg_val, 0, static_cast(1)); + simdgroup_offsets[simd_lane_id] = exclusive; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // This thread's exclusive offset = inclusive_scan - chunk_total + + // simdgroup_offset + int thread_offset = val - chunk_total + simdgroup_offsets[simd_group_id]; + + // Write block_offsets for this thread's chunk using a serial exclusive scan + int running = thread_offset; + for (uint i = start; i < end; i++) { + block_offsets[i] = running; + running += block_sums[i]; + } + + if (lid == tgsize - 1) { + *total_nonzero = + simdgroup_offsets[num_simds - 1] + simdgroup_totals[num_simds - 1]; + } +} + +// Scatter the multi-dimensional indices of nonzero elements. +// Output layout: out[position * ndim + d] = index along dimension d. +template +[[max_total_threads_per_threadgroup(1024)]] +kernel void scatter_nonzero_indices( + const device T* input [[buffer(0)]], + const device int* prefix [[buffer(1)]], + device int64_t* output [[buffer(2)]], + constant int& ndim [[buffer(3)]], + constant int64_t* sizes [[buffer(4)]], + constant int* block_offsets [[buffer(5)]], + constant int& max_entries [[buffer(6)]], + uint tid [[thread_position_in_grid]], + uint tgid [[threadgroup_position_in_grid]]) { + if (!is_nonzero(input[tid])) + return; + + int pos = block_offsets[tgid] + prefix[tid]; + if (pos >= max_entries) + return; + + uint flat = tid; + for (int d = ndim - 1; d >= 0; d--) { + int64_t dim_size = sizes[d]; + output[pos * ndim + d] = + static_cast(flat % static_cast(dim_size)); + flat /= static_cast(dim_size); + } +} + +#define REGISTER_NONZERO_KERNELS(DTYPE) \ + template [[host_name("count_nonzero_prefix_sum_" #DTYPE)]] [[kernel]] void \ + count_nonzero_prefix_sum( \ + const device DTYPE* input [[buffer(0)]], \ + device int* prefix [[buffer(1)]], \ + device int* block_sums [[buffer(2)]], \ + uint tid [[thread_position_in_grid]], \ + uint lid [[thread_position_in_threadgroup]], \ + uint tgsize [[threads_per_threadgroup]], \ + uint tgid [[threadgroup_position_in_grid]], \ + uint simd_lane_id [[thread_index_in_simdgroup]], \ + uint simd_group_id [[simdgroup_index_in_threadgroup]]); \ + \ + template [[host_name("scatter_nonzero_indices_" #DTYPE)]] [[kernel]] void \ + scatter_nonzero_indices( \ + const device DTYPE* input [[buffer(0)]], \ + const device int* prefix [[buffer(1)]], \ + device int64_t* output [[buffer(2)]], \ + constant int& ndim [[buffer(3)]], \ + constant int64_t* sizes [[buffer(4)]], \ + constant int* block_offsets [[buffer(5)]], \ + constant int& max_entries [[buffer(6)]], \ + uint tid [[thread_position_in_grid]], \ + uint tgid [[threadgroup_position_in_grid]]) + +REGISTER_NONZERO_KERNELS(float); +REGISTER_NONZERO_KERNELS(half); +REGISTER_NONZERO_KERNELS(bfloat); +REGISTER_NONZERO_KERNELS(long); +REGISTER_NONZERO_KERNELS(int); +REGISTER_NONZERO_KERNELS(short); +REGISTER_NONZERO_KERNELS(char); +REGISTER_NONZERO_KERNELS(uchar); +REGISTER_NONZERO_KERNELS(bool); +REGISTER_NONZERO_KERNELS(float2); +REGISTER_NONZERO_KERNELS(half2); diff --git a/aten/src/ATen/native/mps/kernels/SamplingHelpers.h b/aten/src/ATen/native/mps/kernels/SamplingHelpers.h new file mode 100644 index 0000000000000..2fd06dfc26247 --- /dev/null +++ b/aten/src/ATen/native/mps/kernels/SamplingHelpers.h @@ -0,0 +1,42 @@ +#pragma once + +// Shared cubic interpolation helpers used by both GridSampler and UpSample +// kernels. Based on the Catmull-Rom spline with A=-0.75. +// See +// https://en.wikipedia.org/wiki/Bicubic_interpolation#Bicubic_convolution_algorithm + +template +accscalar_t cubic_convolution1(accscalar_t x, accscalar_t A) { + return ((A + 2) * x - (A + 3)) * x * x + 1; +} + +template +accscalar_t cubic_convolution2(accscalar_t x, accscalar_t A) { + return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A; +} + +template +void get_cubic_coefficients(accscalar_t coeffs[4], accscalar_t t) { + accscalar_t A = -0.75; + + accscalar_t x1 = t; + coeffs[0] = cubic_convolution2(x1 + 1.0, A); + coeffs[1] = cubic_convolution1(x1, A); + + accscalar_t x2 = 1.0 - t; + coeffs[2] = cubic_convolution1(x2, A); + coeffs[3] = cubic_convolution2(x2 + 1.0, A); +} + +template +accscalar_t cubic_interp1d( + scalar_t x0, + scalar_t x1, + scalar_t x2, + scalar_t x3, + accscalar_t t) { + accscalar_t coeffs[4]; + get_cubic_coefficients(coeffs, t); + + return x0 * coeffs[0] + x1 * coeffs[1] + x2 * coeffs[2] + x3 * coeffs[3]; +} diff --git a/aten/src/ATen/native/mps/kernels/ScanKernel.metal b/aten/src/ATen/native/mps/kernels/ScanKernel.metal index de493af7aaa05..1fde656f46e69 100644 --- a/aten/src/ATen/native/mps/kernels/ScanKernel.metal +++ b/aten/src/ATen/native/mps/kernels/ScanKernel.metal @@ -3,26 +3,16 @@ using namespace metal; #include +#include #include -using c10::metal::accum_t; +using namespace c10::metal; struct LogAddExp { template T operator()(T x, T y) { - // Reference: - // https://www.tensorflow.org/api_docs/python/tf/math/cumulative_logsumexp - T min_val = c10::metal::min(x, y); - T max_val = c10::metal::max(x, y); - - if (min_val != max_val || metal::isfinite(min_val)) { - // nan will be propagated here - return c10::metal::log1p(metal::exp(min_val - max_val)) + max_val; - } else { - // special case to correctly handle infinite cases - return x; - } - }; + return c10::metal::logaddexp(x, y); + } }; C10_METAL_CONSTEXPR auto simd_size = c10::metal::simdgroup_size; @@ -123,6 +113,29 @@ struct LogCumSumExpOp { } }; +template > +struct CumProdOp { + static constexpr constant acc_t init = + is_complex_v ? acc_t(1, 0) : static_cast(1); + + acc_t operator()(acc_t a, acc_t b) { + return c10::metal::mul(a, b); + } + + acc_t simd_scan(acc_t x) { + for (int i = 1; i <= 16; i *= 2) { + acc_t other = simd_shuffle_and_fill_up(x, init, i); + x = this->operator()(x, other); + } + return x; + } + + acc_t simd_exclusive_scan(acc_t x) { + x = simd_scan(x); + return simd_shuffle_and_fill_up(x, init, 1); + } +}; + // Pair structure to hold value and index for cummin/cummax operations template > struct ValueIndexPair { @@ -765,6 +778,11 @@ kernel void scan_with_indices_outer_dim( REGISTER_SCAN_OP(logcumsumexp, LogCumSumExpOp, float, 4); REGISTER_SCAN_OP(logcumsumexp, LogCumSumExpOp, half, 4); REGISTER_SCAN_OP(logcumsumexp, LogCumSumExpOp, bfloat, 4); +REGISTER_SCAN_OP(logcumsumexp, LogCumSumExpOp, float2, 2); +REGISTER_SCAN_OP(logcumsumexp, LogCumSumExpOp, half2, 4); + +REGISTER_SCAN_OP(cumprod, CumProdOp, float2, 2); +REGISTER_SCAN_OP(cumprod, CumProdOp, half2, 4); // Scan with indices operations for cummin/cummax REGISTER_SCAN_WITH_INDICES_OP(cummin, CumMinOp, float, 4); diff --git a/aten/src/ATen/native/mps/kernels/UpSample.metal b/aten/src/ATen/native/mps/kernels/UpSample.metal index fa9b5a1bb107d..cd4be4d0fc5bc 100644 --- a/aten/src/ATen/native/mps/kernels/UpSample.metal +++ b/aten/src/ATen/native/mps/kernels/UpSample.metal @@ -1,3 +1,4 @@ +#include #include #include #include @@ -5,45 +6,6 @@ using namespace metal; using namespace c10::metal; -// Based on -// https://en.wikipedia.org/wiki/Bicubic_interpolation#Bicubic_convolution_algorithm -template -accscalar_t cubic_convolution1(accscalar_t x, accscalar_t A) { - return ((A + 2) * x - (A + 3)) * x * x + 1; -} - -template -accscalar_t cubic_convolution2(accscalar_t x, accscalar_t A) { - return ((A * x - 5 * A) * x + 8 * A) * x - 4 * A; -} - -template -void get_cubic_upsampling_coefficients(accscalar_t coeffs[4], accscalar_t t) { - accscalar_t A = -0.75; - - accscalar_t x1 = t; - coeffs[0] = cubic_convolution2(x1 + 1.0, A); - coeffs[1] = cubic_convolution1(x1, A); - - // opposite coefficients - accscalar_t x2 = 1.0 - t; - coeffs[2] = cubic_convolution1(x2, A); - coeffs[3] = cubic_convolution2(x2 + 1.0, A); -} - -template -accscalar_t cubic_interp1d( - scalar_t x0, - scalar_t x1, - scalar_t x2, - scalar_t x3, - accscalar_t t) { - accscalar_t coeffs[4]; - get_cubic_upsampling_coefficients(coeffs, t); - - return x0 * coeffs[0] + x1 * coeffs[1] + x2 * coeffs[2] + x3 * coeffs[3]; -} - template accscalar_t area_pixel_compute_source_index( accscalar_t scale, @@ -720,8 +682,8 @@ kernel void upsample_bicubic2d_backward( float x_coeffs[4]; float y_coeffs[4]; - get_cubic_upsampling_coefficients(x_coeffs, t_x); - get_cubic_upsampling_coefficients(y_coeffs, t_y); + get_cubic_coefficients(x_coeffs, t_x); + get_cubic_coefficients(y_coeffs, t_y); for (int n = 0; n < output_sizes.x; n++) { for (int c = 0; c < output_sizes.y; ++c) { diff --git a/aten/src/ATen/native/mps/operations/Activation.mm b/aten/src/ATen/native/mps/operations/Activation.mm index 802c648c888d5..205b6bb3872c2 100644 --- a/aten/src/ATen/native/mps/operations/Activation.mm +++ b/aten/src/ATen/native/mps/operations/Activation.mm @@ -22,8 +22,6 @@ #include #include #include -#include -#include #include #include #include @@ -35,86 +33,6 @@ namespace at::native { -Tensor relu_mps(const Tensor& self) { - using namespace mps; - using CachedGraph = MPSUnaryCachedGraph; - - bool executeGatherOp = - !(self.is_contiguous(MemoryFormat::Contiguous) || self.is_contiguous(MemoryFormat::ChannelsLast) || - self.is_contiguous(MemoryFormat::ChannelsLast3d)); - Tensor output = at::empty_like(self, executeGatherOp ? MemoryFormat::Contiguous : MemoryFormat::Preserve); - - if (output.numel() == 0) { - return output; - } - - MPSStream* stream = getCurrentMPSStream(); - @autoreleasepool { - std::string key = "relu" + getTensorsStringKey({self}); - auto cachedGraph = LookUpOrCreateCachedGraph(key, [&](auto mpsGraph, auto newCachedGraph) { - MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self); - // passing selector of reLUWithTensor on the mpsGraph object - MPSGraphTensor* outputTensor = [mpsGraph reLUWithTensor:inputTensor name:nil]; - - newCachedGraph->inputTensor_ = inputTensor; - newCachedGraph->outputTensor_ = outputTensor; - }); - - Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self, nil, executeGatherOp); - Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, output, nil, false); - - // Create dictionary of inputs and outputs - auto feeds = dictionaryFromPlaceholders(selfPlaceholder); - runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder); - } - - return output; -} - -Tensor& relu_mps_(Tensor& self) { - using namespace mps; - using CachedGraph = MPSUnaryCachedGraph; - - if (self.numel() == 0) { - return self; - } - // Inplace relu - Tensor& output = self; - bool executeGatherOp = - !(self.is_contiguous(MemoryFormat::Contiguous) || self.is_contiguous(MemoryFormat::ChannelsLast) || - self.is_contiguous(MemoryFormat::ChannelsLast3d)); - Tensor out; - if (executeGatherOp) { - out = at::empty_like(self, MemoryFormat::Contiguous); - } - - MPSStream* stream = getCurrentMPSStream(); - - @autoreleasepool { - std::string key = "relu_" + getTensorsStringKey({self}); - auto cachedGraph = LookUpOrCreateCachedGraph(key, [&](auto mpsGraph, auto newCachedGraph) { - MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self); - // passing selector of reLUWithTensor on the mpsGraph object - MPSGraphTensor* outputTensor = [mpsGraph reLUWithTensor:inputTensor name:nil]; - - newCachedGraph->inputTensor_ = inputTensor; - newCachedGraph->outputTensor_ = outputTensor; - }); - - Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self, nil, executeGatherOp); - Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, executeGatherOp ? out : output, nil, false); - - // Create dictionary of inputs and outputs - auto feeds = dictionaryFromPlaceholders(selfPlaceholder); - runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder); - if (executeGatherOp) { - output.copy_(out); - } - } - - return output; -} - TORCH_IMPL_FUNC(log_softmax_mps_out) (const Tensor& self, const int64_t dim, const bool half_to_float, const Tensor& out) { TORCH_CHECK_NOT_IMPLEMENTED(self.scalar_type() != kLong, "MPS doesn't know how to do exponent_i64"); @@ -1216,112 +1134,6 @@ Tensor prelu_mps(const Tensor& self, const Tensor& weight_) { return std::tuple{grad_input, weight_grad}; } -TORCH_IMPL_FUNC(silu_out_mps)(const Tensor& self, const Tensor& result) { - using namespace mps; - using CachedGraph = MPSUnaryCachedGraph; - - TORCH_CHECK(self.is_mps()); - TORCH_CHECK_NOT_IMPLEMENTED(self.scalar_type() != kLong, "MPS doesn't know how to do exponent_i64"); - - // Empty output - if (result.numel() == 0) - return; - - MPSStream* stream = getCurrentMPSStream(); - - bool executeGatherOp = - !(self.is_contiguous(MemoryFormat::Contiguous) || self.is_contiguous(MemoryFormat::ChannelsLast) || - self.is_contiguous(MemoryFormat::ChannelsLast3d)); - Tensor result_ = at::empty_like(self, executeGatherOp ? MemoryFormat::Contiguous : MemoryFormat::Preserve); - - @autoreleasepool { - std::string key = "silu_out_mps:" + getTensorsStringKey({self}); - - auto cachedGraph = LookUpOrCreateCachedGraph(key, [&](auto mpsGraph, auto newCachedGraph) { - MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self); - - MPSGraphTensor* unitTensor = [mpsGraph constantWithScalar:1.0 shape:@[ @1 ] dataType:getMPSDataType(self)]; - MPSGraphTensor* negativeInput = [mpsGraph negativeWithTensor:inputTensor name:nil]; - MPSGraphTensor* expNegativeTensor = [mpsGraph exponentWithTensor:negativeInput name:nil]; - MPSGraphTensor* expPlusOneTensor = [mpsGraph additionWithPrimaryTensor:expNegativeTensor - secondaryTensor:unitTensor - name:nil]; - MPSGraphTensor* outputTensor = [mpsGraph divisionWithPrimaryTensor:inputTensor - secondaryTensor:expPlusOneTensor - name:nil]; - - newCachedGraph->inputTensor_ = inputTensor; - newCachedGraph->outputTensor_ = outputTensor; - }); - - Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self, nil, executeGatherOp); - Placeholder outputPlaceholder = - Placeholder(cachedGraph->outputTensor_, executeGatherOp ? result_ : result, nil, false); - - auto feeds = dictionaryFromPlaceholders(selfPlaceholder); - runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder); - } - if (executeGatherOp) { - result.copy_(result_); - } -} - -TORCH_IMPL_FUNC(silu_backward_out_mps) -(const Tensor& grad_output, const Tensor& self, const Tensor& grad_input) { - using namespace mps; - using CachedGraph = MPSUnaryGradCachedGraph; - TORCH_CHECK(grad_output.is_mps()); - - // Empty output - if (grad_input.numel() == 0) - return; - - MPSStream* stream = getCurrentMPSStream(); - - @autoreleasepool { - std::string key = "silu_out_backward_mps:" + getTensorsStringKey({grad_output}); - - auto cachedGraph = LookUpOrCreateCachedGraph(key, [&](auto mpsGraph, auto newCachedGraph) { - MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self); - MPSGraphTensor* gradOutputTensor = mpsGraphRankedPlaceHolder(mpsGraph, grad_output); - - MPSGraphTensor* unitTensor = [mpsGraph constantWithScalar:1.0 shape:@[ @1 ] dataType:getMPSDataType(grad_output)]; - MPSGraphTensor* negativeInput = [mpsGraph negativeWithTensor:inputTensor name:nil]; - MPSGraphTensor* expNegativeTensor = [mpsGraph exponentWithTensor:negativeInput name:nil]; - MPSGraphTensor* expPlusOneTensor = [mpsGraph additionWithPrimaryTensor:expNegativeTensor - secondaryTensor:unitTensor - name:nil]; - MPSGraphTensor* sigmoidTensor = [mpsGraph reciprocalWithTensor:expPlusOneTensor name:nil]; - MPSGraphTensor* oneMinusSigmoid = [mpsGraph subtractionWithPrimaryTensor:unitTensor - secondaryTensor:sigmoidTensor - name:nil]; - MPSGraphTensor* inputTimesDiff = [mpsGraph multiplicationWithPrimaryTensor:inputTensor - secondaryTensor:oneMinusSigmoid - name:nil]; - MPSGraphTensor* onePlusTensor = [mpsGraph additionWithPrimaryTensor:unitTensor - secondaryTensor:inputTimesDiff - name:nil]; - MPSGraphTensor* gradTensor = [mpsGraph multiplicationWithPrimaryTensor:sigmoidTensor - secondaryTensor:onePlusTensor - name:nil]; - MPSGraphTensor* gradInputTensor = [mpsGraph multiplicationWithPrimaryTensor:gradTensor - secondaryTensor:gradOutputTensor - name:nil]; - - newCachedGraph->gradOutputTensor_ = gradOutputTensor; - newCachedGraph->inputTensor_ = inputTensor; - newCachedGraph->gradInputTensor_ = gradInputTensor; - }); - - Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self); - Placeholder gradOutputPlaceholder = Placeholder(cachedGraph->gradOutputTensor_, grad_output); - Placeholder gradInputPlaceholder = Placeholder(cachedGraph->gradInputTensor_, grad_input); - - auto feeds = dictionaryFromPlaceholders(selfPlaceholder, gradOutputPlaceholder); - runMPSGraph(stream, cachedGraph->graph(), feeds, gradInputPlaceholder); - } -} - // ------------------------------------------------- // Hardtanh backward diff --git a/aten/src/ATen/native/mps/operations/ActivationKernel.mm b/aten/src/ATen/native/mps/operations/ActivationKernel.mm index f6d3ad986ade0..2138b0ddb6a96 100644 --- a/aten/src/ATen/native/mps/operations/ActivationKernel.mm +++ b/aten/src/ATen/native/mps/operations/ActivationKernel.mm @@ -4,6 +4,18 @@ #include #include #include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +#include +#include +#include +#include +#include +#include +#include +#endif #include #include @@ -15,6 +27,25 @@ #include #endif +Tensor relu_mps(const Tensor& self) { + TORCH_CHECK(!self.is_complex(), "relu is not supported for complex types"); + auto output = at::empty_like(self); + if (output.numel() == 0) + return output; + auto iter = at::TensorIteratorConfig().add_output(output).add_input(self).build(); + lib.exec_unary_kernel(iter, "relu", /*alpha=*/std::nullopt, /*scalar_arg_type=*/std::nullopt, /*supports_vec4=*/true); + return output; +} + +Tensor& relu_mps_(Tensor& self) { + TORCH_CHECK(!self.is_complex(), "relu is not supported for complex types"); + if (self.numel() == 0) + return self; + auto iter = at::TensorIteratorConfig().add_output(self).add_input(self).set_check_mem_overlap(false).build(); + lib.exec_unary_kernel(iter, "relu", /*alpha=*/std::nullopt, /*scalar_arg_type=*/std::nullopt, /*supports_vec4=*/true); + return self; +} + static void hardshrink_kernel(TensorIteratorBase& iter, const Scalar& lambda = 0.5) { lib.exec_unary_kernel(iter, "hardshrink", lambda); } @@ -67,6 +98,30 @@ static void elu_backward_kernel(TensorIteratorBase& iter, }); } +static void silu_kernel(TensorIteratorBase& iter) { + if (isComplexType(iter.common_dtype())) { + auto out = iter.output(0); + auto self = iter.input(0); + at::mul_out(out, self, at::sigmoid(self)); + return; + } + lib.exec_unary_kernel(iter, "silu", /*alpha=*/std::nullopt, /*scalar_arg_type=*/std::nullopt, /*supports_vec4=*/true); +} + +static void silu_backward_kernel(TensorIteratorBase& iter) { + if (isComplexType(iter.common_dtype())) { + auto grad_input = iter.output(0); + auto grad_output = iter.input(0); + auto self = iter.input(1); + auto sig = at::sigmoid(self); + auto one_minus_sig = at::rsub(sig, 1); + auto inner = at::add(at::mul(self, one_minus_sig), 1); + grad_input.copy_(at::mul(grad_output, at::mul(sig, inner))); + return; + } + lib.exec_binary_kernel(iter, "silu_backward"); +} + static void leaky_relu_kernel(TensorIteratorBase& iter, const Scalar& negative_slope) { lib.exec_unary_kernel(iter, "leaky_relu", negative_slope); } @@ -86,5 +141,7 @@ static void leaky_relu_backward_kernel(TensorIteratorBase& iter, const Scalar& n REGISTER_DISPATCH(elu_backward_stub, elu_backward_kernel); REGISTER_DISPATCH(leaky_relu_stub, leaky_relu_kernel); REGISTER_DISPATCH(leaky_relu_backward_stub, leaky_relu_backward_kernel); +REGISTER_DISPATCH(silu_stub, silu_kernel); +REGISTER_DISPATCH(silu_backward_stub, silu_backward_kernel); } // namespace at::native diff --git a/aten/src/ATen/native/mps/operations/Attention.mm b/aten/src/ATen/native/mps/operations/Attention.mm index 90600cd8a7b85..696094d800320 100644 --- a/aten/src/ATen/native/mps/operations/Attention.mm +++ b/aten/src/ATen/native/mps/operations/Attention.mm @@ -439,13 +439,42 @@ } std::tuple _scaled_dot_product_attention_math_mps(const Tensor& query, - const Tensor& key, - const Tensor& value, + const Tensor& key_, + const Tensor& value_, const std::optional& attn_mask, double dropout_p, bool is_causal, const std::optional& dropout_mask, - std::optional scale) { + std::optional scale, + bool enable_gqa) { + TORCH_CHECK_NOT_IMPLEMENTED(c10::isFloatingType(query.scalar_type()), + "scaled_dot_product_attention for MPS does not support dtype ", + query.scalar_type()); + TORCH_CHECK_NOT_IMPLEMENTED(c10::isFloatingType(key_.scalar_type()), + "scaled_dot_product_attention for MPS does not support dtype ", + key_.scalar_type()); + TORCH_CHECK_NOT_IMPLEMENTED(c10::isFloatingType(value_.scalar_type()), + "scaled_dot_product_attention for MPS does not support dtype ", + value_.scalar_type()); + const auto any_nested = query.is_nested() || key_.is_nested() || value_.is_nested(); + const auto all_contiguous = + query.is_contiguous_or_false() && key_.is_contiguous_or_false() && value_.is_contiguous_or_false(); + auto key = key_; + auto value = value_; + if (enable_gqa) { + int64_t q_heads = query.size(-3); + int64_t k_heads = key_.size(-3); + int64_t repeat_factor = q_heads / k_heads; + + if (repeat_factor > 1) { + TORCH_CHECK(q_heads % k_heads == 0, + "For GQA, the query tensor's head dimension (" + std::to_string(q_heads) + + ") must be divisible by the key tensor's head dimension (" + std::to_string(k_heads) + ")."); + key = key_.repeat_interleave(repeat_factor, /*dim=*/-3); + value = value_.repeat_interleave(repeat_factor, /*dim=*/-3); + } + } + auto query_tuple = ensure_4d(query); Tensor q_ = std::get<0>(query_tuple); bool unsqueezed = std::get<1>(query_tuple); diff --git a/aten/src/ATen/native/mps/operations/BinaryKernel.mm b/aten/src/ATen/native/mps/operations/BinaryKernel.mm index 66190e185f3dc..9f78ab11dfab1 100644 --- a/aten/src/ATen/native/mps/operations/BinaryKernel.mm +++ b/aten/src/ATen/native/mps/operations/BinaryKernel.mm @@ -190,6 +190,68 @@ static void lerp_scalar_mps_kernel(at::TensorIteratorBase& iter, const Scalar& w lib.exec_binary_kernel(iter, "lerp_alpha", weight); } +static void lerp_tensor_mps_kernel(at::TensorIteratorBase& iter) { + using namespace mps; + auto type_str = scalarToMetalTypeString(iter.common_dtype()); + auto numel = static_cast(iter.numel()); + auto ndim = static_cast(iter.ndim()); + + // simple elementwise kernel for dense tensors + if (iter.is_contiguous()) { + auto pso = lib.getPipelineStateForFunc("lerp_tensor_dense_" + type_str); + dispatch_sync_with_rethrow(getCurrentMPSStream()->queue(), ^() { + auto computeEncoder = getCurrentMPSStream()->commandEncoder(); + [computeEncoder setComputePipelineState:pso]; + bind_iter_tensors(computeEncoder, iter); + mtl_dispatch1DJob(computeEncoder, pso, numel); + }); + return; + } + + // Scalar weight broadcast path + if (ndim == 1 && iter.strides(3)[0] == 0) { + auto pso = lib.getPipelineStateForFunc("lerp_tensor_scalar_weight_" + type_str); + dispatch_sync_with_rethrow(getCurrentMPSStream()->queue(), ^() { + auto computeEncoder = getCurrentMPSStream()->commandEncoder(); + [computeEncoder setComputePipelineState:pso]; + bind_iter_tensors(computeEncoder, iter); + mtl_dispatch1DJob(computeEncoder, pso, numel); + }); + return; + } + + // 2D/3D: multi-dimensional dispatch, to avoid integer division for coordinates + if (ndim >= 2 && ndim <= 3) { + auto pso = lib.getPipelineStateForFunc(fmt::format("lerp_tensor_strided_{}d_{}", ndim, type_str)); + dispatch_sync_with_rethrow(getCurrentMPSStream()->queue(), ^() { + auto computeEncoder = getCurrentMPSStream()->commandEncoder(); + [computeEncoder setComputePipelineState:pso]; + bind_iter_tensors(computeEncoder, iter); + mtl_setArgs<4>(computeEncoder, iter.strides(0), iter.strides(1), iter.strides(2), iter.strides(3)); + auto sizes = iter.shape(); + auto maxTg = [pso maxTotalThreadsPerThreadgroup]; + auto tg_x = std::min(static_cast(sizes[0]), maxTg); + auto tg_y = std::min(static_cast(sizes[1]), maxTg / tg_x); + auto grid_z = ndim > 2 ? static_cast(sizes[2]) : 1; + auto tg_z = std::min(grid_z, std::max(maxTg / (tg_x * tg_y), (NSUInteger)1)); + [computeEncoder dispatchThreads:MTLSizeMake(sizes[0], sizes[1], grid_z) + threadsPerThreadgroup:MTLSizeMake(tg_x, tg_y, tg_z)]; + }); + return; + } + + // General strided fallback + auto pso = lib.getPipelineStateForFunc("lerp_tensor_strided_" + type_str); + dispatch_sync_with_rethrow(getCurrentMPSStream()->queue(), ^() { + auto computeEncoder = getCurrentMPSStream()->commandEncoder(); + [computeEncoder setComputePipelineState:pso]; + bind_iter_tensors(computeEncoder, iter); + mtl_setArgs<4>( + computeEncoder, iter.shape(), iter.strides(0), iter.strides(1), iter.strides(2), iter.strides(3), ndim); + mtl_dispatch1DJob(computeEncoder, pso, numel); + }); +} + static void native_dropout_mask_and_scale_mps_kernel(at::TensorIteratorBase& iter, const Scalar& scale) { lib.exec_binary_kernel(iter, "native_dropout_mask_and_scale", scale); } @@ -261,6 +323,7 @@ static void gcd_mps_kernel(TensorIteratorBase& iter) { REGISTER_DISPATCH(polar_stub, &polar_mps_kernel); REGISTER_DISPATCH(complex_stub, &complex_mps_kernel); REGISTER_DISPATCH(lerp_kernel_scalar_weight, &lerp_scalar_mps_kernel) +REGISTER_DISPATCH(lerp_kernel_tensor_weight, &lerp_tensor_mps_kernel) REGISTER_DISPATCH(mul_stub, &mul_mps_kernel) REGISTER_DISPATCH(div_true_stub, &div_true_mps_kernel) REGISTER_DISPATCH(div_floor_stub, &div_floor_mps_kernel) diff --git a/aten/src/ATen/native/mps/operations/ConstantOps.mm b/aten/src/ATen/native/mps/operations/ConstantOps.mm index e36ac4dc45246..c0d3b795ded9f 100644 --- a/aten/src/ATen/native/mps/operations/ConstantOps.mm +++ b/aten/src/ATen/native/mps/operations/ConstantOps.mm @@ -1,120 +1,89 @@ // Copyright © 2022 Apple Inc. #define TORCH_ASSERT_ONLY_METHOD_OPERATORS +#include #include +#include #ifndef AT_PER_OPERATOR_HEADERS #include #include #else -#include -#include -#include +#include #endif namespace at::native { -static Tensor& fill_scalar_mps_impl(Tensor& self, const Scalar& value) { - using namespace mps; +#ifndef PYTORCH_JIT_COMPILE_SHADERS +static auto& lib = mps::MetalShaderLibrary::getBundledLibrary(); +#else +#include +#endif - if (self.numel() == 0) { - return self; - } - Tensor output = self; - bool needsCopyToOutput = false; - if (needsGather(self)) { - output = at::empty(self.sizes(), self.scalar_type(), std::nullopt, kMPS, std::nullopt, std::nullopt); - needsCopyToOutput = true; +static void fill_mps_kernel(TensorIterator& iter, const Scalar& value) { + using namespace mps; + if (iter.numel() == 0) { + return; } - struct CachedGraph : public MPSCachedGraph { - CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {} - MPSGraphTensor* inputTensor_ = nil; - MPSGraphTensor* outputTensor_ = nil; - }; - - @autoreleasepool { - std::string key = "fill_scalar_mps_impl" + getTensorsStringKey(self) + ":" + std::to_string(value.toDouble()); - - auto cachedGraph = LookUpOrCreateCachedGraph(key, [&](auto mpsGraph, auto newCachedGraph) { - MPSGraphTensor* inputTensor = mpsGraphScalarPlaceHolder(mpsGraph, getMPSDataType(self.scalar_type())); - MPSGraphTensor* outputTensor = [mpsGraph identityWithTensor:inputTensor name:nil]; - newCachedGraph->inputTensor_ = inputTensor; - newCachedGraph->outputTensor_ = outputTensor; - }); - - auto mpsScalar = getMPSScalar(value, self.scalar_type()); - auto mpsScalarData = getMPSGraphTensorFromScalar(getCurrentMPSStream(), mpsScalar); - NSDictionary* feeds = @{cachedGraph->inputTensor_ : mpsScalarData}; - - Placeholder outputPlaceholder = - Placeholder(cachedGraph->outputTensor_, needsCopyToOutput ? output : self, nullptr, !needsCopyToOutput); - - NSDictionary* results = - @{outputPlaceholder.getMPSGraphTensor() : outputPlaceholder.getMPSGraphTensorData()}; - - runMPSGraph(getCurrentMPSStream(), cachedGraph->graph(), feeds, results); - - if (needsCopyToOutput) { - self.copy_(output); + // Metal compute kernels use uint (32-bit) thread indices; decompose large + // tensors into chunks that fit in 32-bit indexing. + if (!iter.can_use_32bit_indexing()) { + for (auto&& sub_iter : iter.with_32bit_indexing()) { + fill_mps_kernel(sub_iter, value); } + return; } - return self; -} - -static Tensor& fill_mps_tensor_(Tensor& self, uint8_t value) { - TORCH_INTERNAL_ASSERT(self.is_contiguous()); + const Tensor& self = iter.tensor(0); + const auto dtype = self.scalar_type(); const auto stream = getCurrentMPSStream(); - auto storage_byte_offset = self.storage_offset() * self.itemsize(); - stream->fill(mps::getMTLBufferStorage(self), value, self.nbytes(), storage_byte_offset); - return self; -} - -Tensor& fill_scalar_mps(Tensor& self, const Scalar& value) { - if (isComplexType(self.scalar_type())) { - auto self_as_real = at::view_as_real(self); - auto self_as_real_real = self_as_real.select(self.dim(), 0); - auto self_as_real_imag = self_as_real.select(self.dim(), 1); - if (value.isComplex()) { - auto value_cdouble = value.to>(); - fill_scalar_mps_impl(self_as_real_real, value_cdouble.real()); - fill_scalar_mps_impl(self_as_real_imag, value_cdouble.imag()); - return self; - } - fill_scalar_mps_impl(self_as_real_real, value); - fill_scalar_mps_impl(self_as_real_imag, 0.0f); - return self; - } - // check if it's possible to use fillBuffer() to fill the Tensor's storage - if (self.is_contiguous()) { - if (value.toDouble() == 0.0) { - return fill_mps_tensor_(self, 0); - } - if (self.scalar_type() == kBool) { - return fill_mps_tensor_(self, value.toBool()); - } - if (self.scalar_type() == kByte) { - return fill_mps_tensor_(self, value.toByte()); - } - if (self.scalar_type() == kChar) { - return fill_mps_tensor_(self, value.toChar()); - } + const auto type_str = scalarToMetalTypeString(dtype); + const bool can_fill_linearly = self.is_non_overlapping_and_dense(); + + // For tensors with gaps or overlaps (e.g. stride-2 slices) use a 2D strided + // kernel: tid.y indexes dim 0 directly (no division), tid.x is the linear + // index for the remaining dims. Consecutive threads in x write consecutive + // addresses in the innermost dimension, giving coalesced writes. + if (!can_fill_linearly) { + auto fillPSO = lib.getPipelineStateForFunc(fmt::format("fill_scalar_strided_{}", type_str)); + const int64_t dim0_size = iter.ndim() > 0 ? iter.shape()[0] : 1; + const int64_t inner_numel = iter.numel() / dim0_size; + const uint32_t ndim = static_cast(iter.ndim()); + dispatch_sync_with_rethrow(stream->queue(), ^() { + @autoreleasepool { + auto computeEncoder = stream->commandEncoder(); + auto mpsScalar = getMPSScalar(value, dtype); + [computeEncoder setComputePipelineState:fillPSO]; + bind_iter_tensors(computeEncoder, iter); + mtl_setArgs<1>(computeEncoder, mpsScalar, iter.shape(), iter.strides(0), ndim); + const NSUInteger maxTG = fillPSO.maxTotalThreadsPerThreadgroup; + const MTLSize tgSize = MTLSizeMake(std::min(maxTG, (NSUInteger)inner_numel), 1, 1); + const MTLSize gridSize = MTLSizeMake(inner_numel, dim0_size, 1); + [computeEncoder dispatchThreads:gridSize threadsPerThreadgroup:tgSize]; + } + }); + return; } - return fill_scalar_mps_impl(self, value); -} - -Tensor& fill_tensor_mps_(Tensor& self, const Tensor& value) { - TORCH_CHECK(value.dim() == 0, - "fill_ only supports 0-dimension value tensor but got tensor with ", - value.dim(), - " dimensions."); - Scalar scalar_value = value.item(); - return fill_scalar_mps(self, scalar_value); + // Single-byte dtypes (bool, uint8, int8) use vec4 kernels that fill + // 4 elements per thread; all others fill 1 element per thread. + const bool is_byte_type = self.element_size() == 1; + const uint32_t numel = static_cast(iter.numel()); + const int64_t threads = is_byte_type ? (numel + 3) / 4 : numel; + + auto fillPSO = lib.getPipelineStateForFunc(fmt::format("fill_scalar_dense_{}", type_str)); + dispatch_sync_with_rethrow(stream->queue(), ^() { + @autoreleasepool { + auto computeEncoder = stream->commandEncoder(); + auto mpsScalar = getMPSScalar(value, dtype); + [computeEncoder setComputePipelineState:fillPSO]; + bind_iter_tensors(computeEncoder, iter); + mtl_setArgs<1>(computeEncoder, mpsScalar, numel); + mtl_dispatch1DJob(computeEncoder, fillPSO, threads); + } + }); } -Tensor& zero_mps_(Tensor& self) { - return fill_scalar_mps(self, 0.0f); -} +REGISTER_DISPATCH(fill_stub, &fill_mps_kernel); } // namespace at::native diff --git a/aten/src/ATen/native/mps/operations/Distributions.mm b/aten/src/ATen/native/mps/operations/Distributions.mm index d1d7bd60f7a16..a1016fe5cdc90 100644 --- a/aten/src/ATen/native/mps/operations/Distributions.mm +++ b/aten/src/ATen/native/mps/operations/Distributions.mm @@ -11,6 +11,8 @@ #include #include #else +#include +#include #include #include #include @@ -419,38 +421,13 @@ Tensor normal_mps(const Tensor& mean, const Tensor& std, std::optional gen) { - TORCH_CHECK(lambda > 0.0, "exponential_ expects lambda > 0.0, but found lambda=", lambda); - - mps::RandomOpBlock random_op_block = ^RandomOpFn(cachedGraph, randomTensor) { - MPSGraph* mpsGraph = cachedGraph->graph(); - MPSGraphTensor* unitTensor = [mpsGraph constantWithScalar:1.0f dataType:randomTensor.dataType]; - MPSGraphTensor* minusLambdaTensor = [mpsGraph constantWithScalar:-lambda dataType:randomTensor.dataType]; - MPSGraphTensor* subtractTensor = [mpsGraph subtractionWithPrimaryTensor:unitTensor - secondaryTensor:randomTensor - name:nil]; - MPSGraphTensor* logTensor = [mpsGraph logarithmWithTensor:subtractTensor name:nil]; - return [mpsGraph divisionWithPrimaryTensor:logTensor secondaryTensor:minusLambdaTensor name:nil]; - }; - auto eps = std::numeric_limits::epsilon(); - return mps::random_mps_impl(self, - eps, - 1.0, - std::nullopt, - std::nullopt, - MPSGraphRandomDistributionUniform, - gen, - "exponential_mps_:" + std::to_string(lambda), - random_op_block); -} - static Tensor& distribution_kernel_mps_impl(Tensor& self, double param1, double param2, const std::string& kernel_name, int64_t randoms_per_element, - std::optional gen) { + std::optional gen, + int64_t elements_per_thread = 1) { if (self.numel() == 0) { return self; } @@ -483,7 +460,13 @@ Tensor normal_mps(const Tensor& mean, const Tensor& std, std::optional{static_cast(param1), static_cast(param2)}, std::array{seed, base_offset}); - mtl_dispatch1DJob(computeEncoder, pso, output.numel()); + if (elements_per_thread > 1) { + auto numel = static_cast(output.numel()); + mtl_setBytes(computeEncoder, numel, 3); + mtl_dispatch1DJob(computeEncoder, pso, (numel + elements_per_thread - 1) / elements_per_thread); + } else { + mtl_dispatch1DJob(computeEncoder, pso, output.numel()); + } } }); } @@ -495,6 +478,11 @@ Tensor normal_mps(const Tensor& mean, const Tensor& std, std::optional gen) { + TORCH_CHECK(lambda > 0.0, "exponential_ expects lambda > 0.0, but found lambda=", lambda); + return distribution_kernel_mps_impl(self, lambda, 0.0, "exponential", 1, gen, /*elements_per_thread=*/4); +} + Tensor& cauchy_mps_(Tensor& self, double median, double sigma, std::optional gen) { TORCH_CHECK(sigma > 0.0, "cauchy_ expects sigma > 0.0, but found sigma=", sigma); return distribution_kernel_mps_impl(self, median, sigma, "cauchy", 1, gen); @@ -511,6 +499,74 @@ Tensor normal_mps(const Tensor& mean, const Tensor& std, std::optional gen) { + if (alpha.numel() == 0) { + return at::empty_like(alpha); + } + + using namespace mps; + + auto mps_gen = get_generator_or_default(gen, at::mps::detail::getDefaultMPSGenerator()); + auto stream = getCurrentMPSStream(); + Tensor ret = at::empty_like(alpha, alpha.options(), at::MemoryFormat::Contiguous); + auto alpha_contig = alpha.contiguous(); + + @autoreleasepool { + auto pso = lib.getPipelineStateForFunc("standard_gamma_" + scalarToMetalTypeString(ret)); + + int64_t seed; + int64_t base_offset; + // Each thread may consume up to GAMMA_RANDOMS_STRIDE random numbers + constexpr int64_t GAMMA_RANDOMS_STRIDE = 32; + { + // See Note [Acquire lock when using random generators] + std::lock_guard lock(mps_gen->mutex_); + seed = static_cast(mps_gen->current_seed()); + base_offset = static_cast(mps_gen->get_offset()); + mps_gen->set_offset(base_offset + GAMMA_RANDOMS_STRIDE * ret.numel()); + } + + dispatch_sync_with_rethrow(stream->queue(), ^() { + @autoreleasepool { + auto computeEncoder = stream->commandEncoder(); + [computeEncoder setComputePipelineState:pso]; + mtl_setArgs(computeEncoder, ret, alpha_contig, std::array{seed, base_offset}); + mtl_dispatch1DJob(computeEncoder, pso, ret.numel()); + } + }); + } + + return ret; +} + +Tensor _standard_gamma_grad_mps(const Tensor& self, const Tensor& output) { + if (self.numel() == 0) { + return at::empty_like(self); + } + + using namespace mps; + + auto stream = getCurrentMPSStream(); + Tensor ret = at::empty_like(self, self.options(), at::MemoryFormat::Contiguous); + const auto self_contig = self.contiguous(); + const auto output_contig = output.contiguous(); + + @autoreleasepool { + auto pso = lib.getPipelineStateForFunc("standard_gamma_grad_" + scalarToMetalTypeString(ret)); + + dispatch_sync_with_rethrow(stream->queue(), ^() { + @autoreleasepool { + auto computeEncoder = stream->commandEncoder(); + [computeEncoder setComputePipelineState:pso]; + mtl_setArgs(computeEncoder, ret, self_contig, output_contig); + mtl_dispatch1DJob(computeEncoder, pso, ret.numel()); + } + }); + } + + return ret; +} + Tensor& randperm_out_mps(int64_t n, std::optional generator, Tensor& result) { TORCH_CHECK(n >= 0, "n must be non-negative, got", n); TORCH_CHECK(!generator.has_value() || (generator.has_value() && result.device() == generator->device()), diff --git a/aten/src/ATen/native/mps/operations/Eye.mm b/aten/src/ATen/native/mps/operations/Eye.mm index 592654dd57740..1b603304f565e 100644 --- a/aten/src/ATen/native/mps/operations/Eye.mm +++ b/aten/src/ATen/native/mps/operations/Eye.mm @@ -1,5 +1,7 @@ #define TORCH_ASSERT_ONLY_METHOD_OPERATORS +#include #include + #ifndef AT_PER_OPERATOR_HEADERS #include #include @@ -7,95 +9,76 @@ #include #endif -// Steps to add op for MPS backend: -// 1. Register the op in aten/src/ATen/native/native_functions.yaml with the "MPS" dispatch key -// 2. Define the function interface for the MPS backend similar to other -// backends depending on whether its structured or non-structured -// 3. Add boiler-plate error checking code as expected for the Op -// 4. The code structure roughly follows the pattern -// a) get the MPS stream handle to encode work onto -// b) get an instance of MPSGraphCache and create a key unique to the Graph -// needed for implementing this Op. Any shape, dataType or parameter -// passed to the MPSGraph during its construction will need to be included -// here. -// c) Create the graph using make_mps_graph() and add operations to the -// instance of MPSGraph. This is if the Cache->lookup() fails. -// d) Store the MPSGraphTensors for inputs and output which are needed at -// runtime. -// e) Use the CachedGraph instance's inputs and output to create Placeholders -// You will need to pass in Tensor to create MPSGraphTensorData objects. -// f) Using MPSGraphTensor and MPSGraphTensorData instances create a feeds -// dictionary. -// g) Then call runMPSGraph() with input params and return the result. -// - namespace at::native { +using namespace mps; + +#ifndef PYTORCH_JIT_COMPILE_SHADERS +static auto& lib = MetalShaderLibrary::getBundledLibrary(); +#else +#include +#endif Tensor& eye_out_mps(int64_t n, Tensor& result) { // the default value of `m` equals to `n` return eye_out_mps(n, n, result); } -using namespace mps; - Tensor& eye_out_mps(int64_t n, int64_t m, Tensor& result) { - // This is one example of boiler-plate error checking, taking after CPU/CUDA counterparts TORCH_CHECK(n >= 0, "n must be greater or equal to 0, got ", n); TORCH_CHECK(m >= 0, "m must be greater or equal to 0, got ", m); result.resize_({n, m}); - result.zero_(); - // Handle empty outputs if (result.numel() == 0) return result; - // Get MPS stream - MPSStream* stream = getCurrentMPSStream(); - - auto outputDataType = result.scalar_type(); - // Derive from MPSCachedGraph - // This structure is used to cache an MPSGraph with certain keys, so that we don't have to compile the same MPSGraph - // time and time again for the same operation The keys of this structure are based on the inputs and outputs needed - // for the operation here, we don't have any input tensors, just an output tensor. - // If the operator to be added is unary or binary, instead of creating a new CachedGraph struct yourself, please - // consider using `MPSUnaryCachedGraph` or `MPSBinaryCachedGraph` and their corresponding Grad versions in - // `OperationUtils.h`. - struct CachedGraph : public MPSCachedGraph { - CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {} - MPSGraphTensor* outputTensor_ = nil; - }; - - @autoreleasepool { - // A key is used to identify the MPSGraph which was created once, and can be reused if the parameters, data types - // etc match the earlier created MPSGraph - std::string key = "eye_out_mps:" + getTensorsStringKey({result}); - auto cachedGraph = LookUpOrCreateCachedGraph(key, [&](auto* mpsGraph, auto* newCachedGraph) { - MPSGraphTensor* onesTensor = [mpsGraph constantWithScalar:1.0f - shape:getMPSShape(result) - dataType:getMPSDataType(outputDataType)]; - - // Here we can call the MPSGraph API needed to execute the operation. - // The API details can be found here: - // https://developer.apple.com/documentation/metalperformanceshadersgraph/mpsgraph - MPSGraphTensor* outputTensor = [mpsGraph bandPartWithTensor:onesTensor numLower:0 numUpper:0 name:nil]; - - if ([outputTensor dataType] != getMPSDataType(outputDataType)) { - outputTensor = castMPSTensor(mpsGraph, outputTensor, outputDataType); + MPSStream* mpsStream = getCurrentMPSStream(); + auto stride0 = result.stride(0); + auto stride1 = result.stride(1); + + // Small tensors: single-pass 2D kernel (one dispatch, no zero_() overhead). + // Large tensors: zero_() + diagonal fill (memset is faster than n*m branching writes). + constexpr int64_t kSinglePassThreshold = 1 << 22; + + if (n * m <= kSinglePassThreshold) { + auto key = "eye_" + scalarToMetalTypeString(result); + id computeEncoder = mpsStream->commandEncoder(); + id pso = lib.getPipelineStateForFunc(key); + + // Map x to the smaller stride for coalesced writes + bool swap = stride0 < stride1; + auto x_stride = swap ? stride0 : stride1; + auto y_stride = swap ? stride1 : stride0; + auto grid_x = swap ? static_cast(n) : static_cast(m); + auto grid_y = swap ? static_cast(m) : static_cast(n); + + dispatch_sync(mpsStream->queue(), ^() { + @autoreleasepool { + [computeEncoder setComputePipelineState:pso]; + mtl_setArgs(computeEncoder, result, y_stride, x_stride); + + auto maxTg = [pso maxTotalThreadsPerThreadgroup]; + auto tg_x = std::min(grid_x, maxTg); + auto tg_y = std::min(grid_y, std::max(maxTg / tg_x, static_cast(1))); + [computeEncoder dispatchThreads:MTLSizeMake(grid_x, grid_y, 1) + threadsPerThreadgroup:MTLSizeMake(tg_x, tg_y, 1)]; + } + }); + } else { + result.zero_(); + int64_t sz = std::min(n, m); + int64_t diag_stride = stride0 + stride1; + auto key = "eye_diag_" + scalarToMetalTypeString(result); + id computeEncoder = mpsStream->commandEncoder(); + id pso = lib.getPipelineStateForFunc(key); + + dispatch_sync(mpsStream->queue(), ^() { + @autoreleasepool { + [computeEncoder setComputePipelineState:pso]; + mtl_setArgs(computeEncoder, result, diag_stride); + mtl_dispatch1DJob(computeEncoder, pso, sz); } - newCachedGraph->outputTensor_ = outputTensor; }); - - // Create placeholders which use the keys of the CachedGraph to create inputs and outputs of the operation - Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, result); - - // Create dictionary of inputs/feeds and outputs/results - // In this case, there are no inputs, so the feeds are nil - NSDictionary* feeds = nil; - auto results = dictionaryFromPlaceholders(outputPlaceholder); - - // Run the graph - runMPSGraph(stream, cachedGraph->graph(), feeds, results); } return result; diff --git a/aten/src/ATen/native/mps/operations/Indexing.mm b/aten/src/ATen/native/mps/operations/Indexing.mm index dd333d2dc7ebd..1433f4ad9844c 100644 --- a/aten/src/ATen/native/mps/operations/Indexing.mm +++ b/aten/src/ATen/native/mps/operations/Indexing.mm @@ -27,8 +27,6 @@ #include #else #include -#include -#include #include #include #include @@ -42,12 +40,11 @@ #include #include #include +#include #include #include #endif -constexpr auto nonZeroMaxSize = 1UL << 24; - namespace at::native { namespace mps { @@ -292,150 +289,139 @@ static void index_put_kernel_mps(TensorIterator& iter, }); } -static Tensor nonzero_fallback(const Tensor& self) { - return at::nonzero(self.to("cpu")).to("mps"); -} - -static Tensor& nonzero_out_native_mps(const Tensor& self, Tensor& out_) { +// Metal kernel-based nonzero using prefix-sum + scatter. +// Step 1: Per-element exclusive prefix sum of nonzero flags + block totals. +// Step 2: GPU prefix sum of block totals → block offsets + total count. +// Host (optional): Read back total count, allocate output, unless max_element is provided +// Step 3: Scatter multi-dimensional indices into the output. +static void nonzero_impl_mps(const Tensor& self, Tensor& out_, std::optional max_elements) { using namespace mps; - int64_t nDim = self.dim(); + TORCH_CHECK(self.numel() < std::numeric_limits::max(), + "nonzero is not supported for tensors with more than INT_MAX elements, " + "See https://github.com/pytorch/pytorch/issues/51871"); + TORCH_CHECK(out_.dtype() == at::kLong, "Expected output type to be Long, but got ", out_.dtype()); + TORCH_CHECK(self.device() == out_.device(), + "expected self and out to be on the same device, but got out on ", + out_.device(), + " and self on ", + self.device()); + TORCH_CHECK(out_.is_mps()); + + Tensor input = self.contiguous(); + const int64_t nDim = self.dim(); + const auto numel = static_cast(input.numel()); + const auto type_str = scalarToMetalTypeString(input); MPSStream* stream = getCurrentMPSStream(); - using CachedGraph = MPSUnaryCachedGraph; + auto pso_step1 = lib.getPipelineStateForFunc(fmt::format("count_nonzero_prefix_sum_{}", type_str)); + auto pso_step2 = lib.getPipelineStateForFunc("prefix_sum_blocks"); + auto pso_step3 = lib.getPipelineStateForFunc(fmt::format("scatter_nonzero_indices_{}", type_str)); + TORCH_INTERNAL_ASSERT([pso_step1 maxTotalThreadsPerThreadgroup] == [pso_step3 maxTotalThreadsPerThreadgroup], + "nonzero: step 1 and step 3 threadgroup sizes must match"); + + uint32_t threads_per_group = static_cast([pso_step1 maxTotalThreadsPerThreadgroup]); + uint32_t num_blocks = (numel + threads_per_group - 1) / threads_per_group; + + auto tmp = at::empty({input.numel() + 2 * num_blocks + 1}, input.options().dtype(kInt)); + Tensor prefix_buf = tmp.slice(0, 0, numel); + Tensor block_sums_buf = tmp.slice(0, numel, numel + num_blocks); + Tensor block_offsets_buf = tmp.slice(0, numel + num_blocks, numel + 2 * num_blocks); + Tensor total_nonzero_buf = tmp.slice(0, numel + 2 * num_blocks, numel + 2 * num_blocks + 1); + + // Steps 1+2: compute prefix sums and block offsets entirely on GPU dispatch_sync_with_rethrow(stream->queue(), ^() { - stream->synchronize(SyncType::COMMIT_AND_WAIT); + @autoreleasepool { + auto computeEncoder = stream->commandEncoder(); + + [computeEncoder setComputePipelineState:pso_step1]; + mtl_setArgs(computeEncoder, input, prefix_buf, block_sums_buf); + mtl_dispatch1DJob(computeEncoder, pso_step1, numel); + + [computeEncoder setComputePipelineState:pso_step2]; + mtl_setArgs(computeEncoder, block_sums_buf, block_offsets_buf, total_nonzero_buf, num_blocks); + uint32_t tg_size_blocks = std::min(1024u, ((num_blocks + 31) / 32) * 32); + [computeEncoder dispatchThreads:MTLSizeMake(tg_size_blocks, 1, 1) + threadsPerThreadgroup:MTLSizeMake(tg_size_blocks, 1, 1)]; + } }); - int64_t total_nonzero = at::count_nonzero(self).item(); - at::native::resize_output(out_, {total_nonzero, nDim}); - if (out_.numel() == 0) { - return out_; - } - bool contiguous_output = !needsGather(out_); - Tensor out = out_; - if (!contiguous_output) { - out = at::empty_like(out_, MemoryFormat::Contiguous); + if (!max_elements) { + // Dynamic path: sync to learn output size + const int64_t total_nonzero = total_nonzero_buf.item(); + at::native::resize_output(out_, {total_nonzero, nDim}); + max_elements = total_nonzero; } - @autoreleasepool { - std::string key = "nonzero_out_native_mps" + getTensorsStringKey(self); - auto cachedGraph = LookUpOrCreateCachedGraph(key, [&](auto mpsGraph, auto newCachedGraph) { - MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self); + if (out_.numel() == 0) { + return; + } - MPSGraphTensor* outputTensor = [mpsGraph nonZeroIndicesOfTensor:inputTensor name:nil]; + bool contiguous_output = out_.is_contiguous(); + Tensor out = contiguous_output ? out_ : at::empty_like(out_, MemoryFormat::Contiguous); - newCachedGraph->inputTensor_ = inputTensor; - newCachedGraph->outputTensor_ = outputTensor; - }); + int ndim_int = static_cast(nDim); + int max_entries = static_cast(*max_elements); - Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self); - Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, out); - auto feeds = dictionaryFromPlaceholders(selfPlaceholder); - runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder); - } + // Step 3: scatter indices, capped at max_entries + dispatch_sync_with_rethrow(stream->queue(), ^() { + @autoreleasepool { + auto computeEncoder = stream->commandEncoder(); + [computeEncoder setComputePipelineState:pso_step3]; + mtl_setArgs(computeEncoder, input, prefix_buf, out, ndim_int, input.sizes(), block_offsets_buf, max_entries); + mtl_dispatch1DJob(computeEncoder, pso_step3, numel); + } + }); if (!contiguous_output) { out_.copy_(out); } - - return out_; } Tensor& nonzero_out_mps(const Tensor& self, Tensor& out_) { - if (self.is_complex()) { - TORCH_WARN_ONCE("MPS: nonzero op is not supported for complex datatypes. ", - "Falling back on CPU. This may have performance implications."); - Tensor out_fallback = nonzero_fallback(self); - at::native::resize_output(out_, out_fallback.sizes()); - out_.copy_(out_fallback); - return out_; - } - int64_t nDim = self.dim(); if (self.numel() == 0) { at::native::resize_output(out_, {0, nDim}); return out_; } - using namespace mps; - const uint32_t maxDimensions = 16; - - TORCH_CHECK(self.numel() < std::numeric_limits::max(), - "nonzero is not supported for tensors with more than INT_MAX elements, \ - See https://github.com/pytorch/pytorch/issues/51871"); - TORCH_CHECK( - out_.dtype() == at::kLong, "Expected object of scalar type ", at::kLong, " as out, but got ", out_.dtype()); - TORCH_CHECK(self.device() == out_.device(), - "expected self and out to be on the same device, but got out on ", - out_.device(), - " and self on ", - self.device()); - TORCH_CHECK(self.dim() <= maxDimensions, "nonzero is not supported for tensor with more than ", 16, " dimensions"); - TORCH_CHECK(out_.is_mps()); + nonzero_impl_mps(self, out_, std::nullopt); + return out_; +} - if (!is_macos_13_or_newer(MacOSVersion::MACOS_VER_15_0_PLUS) && - (self.numel() >= nonZeroMaxSize || self.is_complex())) { - TORCH_WARN_ONCE("MPS: nonzero op is not natively supported for the provided input on MacOS14", - "Falling back on CPU. This may have performance implications.", - "See github.com/pytorch/pytorch/issues/122916 for further info"); - Tensor out_fallback = nonzero_fallback(self); - at::native::resize_output(out_, out_fallback.sizes()); - out_.copy_(out_fallback); - return out_; - } +Tensor nonzero_mps(const Tensor& self) { + Tensor out = at::empty({0}, self.options().dtype(kLong)); + return nonzero_out_mps(self, out); +} - MPSStream* stream = getCurrentMPSStream(); - using CachedGraph = MPSUnaryCachedGraph; +Tensor& nonzero_static_out_mps(const Tensor& self, int64_t size, int64_t fill_value, Tensor& result) { + TORCH_CHECK(size >= 0, "nonzero_static: 'size' must be an non-negative integer"); - dispatch_sync_with_rethrow(stream->queue(), ^() { - stream->synchronize(SyncType::COMMIT_AND_WAIT); - }); - int64_t total_nonzero = at::count_nonzero(self).item(); - at::native::resize_output(out_, {total_nonzero, nDim}); - if (out_.numel() == 0) { - return out_; + int64_t nDim = self.dim(); + if (result.dim() != 2 || result.size(0) != size || result.size(1) != nDim) { + at::native::resize_output(result, {size, nDim}); } - bool contiguous_output = !needsGather(out_); - Tensor out = out_; - if (!contiguous_output) { - out = at::empty_like(out_, MemoryFormat::Contiguous); + if (result.size(0) == 0 || result.size(1) == 0) { + return result; } - @autoreleasepool { - std::string key = "nonzero_out_native_mps" + getTensorsStringKey(self); - auto cachedGraph = LookUpOrCreateCachedGraph(key, [&](auto mpsGraph, auto newCachedGraph) { - MPSGraphTensor* inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, self); - - MPSGraphTensor* outputTensor = [mpsGraph nonZeroIndicesOfTensor:inputTensor name:nil]; - - newCachedGraph->inputTensor_ = inputTensor; - newCachedGraph->outputTensor_ = outputTensor; - }); + result.fill_(fill_value); - Placeholder selfPlaceholder = Placeholder(cachedGraph->inputTensor_, self); - Placeholder outputPlaceholder = Placeholder(cachedGraph->outputTensor_, out); - auto feeds = dictionaryFromPlaceholders(selfPlaceholder); - runMPSGraph(stream, cachedGraph->graph(), feeds, outputPlaceholder); - } - - if (!contiguous_output) { - out_.copy_(out); + if (self.numel() == 0) { + return result; } - return out_; + nonzero_impl_mps(self, result, size); + return result; } -Tensor nonzero_mps(const Tensor& self) { - if (self.is_complex()) { - TORCH_WARN_ONCE("MPS: nonzero op is not supported for complex datatypes ", - "Falling back on CPU. This may have performance implications."); - return nonzero_fallback(self); - } - - Tensor out = at::empty({0}, self.options().dtype(kLong)); - return nonzero_out_mps(self, out); +Tensor nonzero_static_mps(const Tensor& self, int64_t size, int64_t fill_value) { + TORCH_CHECK(size >= 0, "nonzero_static: 'size' must be an non-negative integer"); + int64_t nDim = self.dim(); + auto result = at::empty({size, nDim}, at::TensorOptions().dtype(kLong).device(kMPS)); + nonzero_static_out_mps(self, size, fill_value, result); + return result; } Tensor masked_select_mps(const Tensor& self, const Tensor& mask) { diff --git a/aten/src/ATen/native/mps/operations/Lerp.mm b/aten/src/ATen/native/mps/operations/Lerp.mm deleted file mode 100644 index 4d13b8b2466bd..0000000000000 --- a/aten/src/ATen/native/mps/operations/Lerp.mm +++ /dev/null @@ -1,49 +0,0 @@ -#define TORCH_ASSERT_ONLY_METHOD_OPERATORS -#include -#include - -#ifndef AT_PER_OPERATOR_HEADERS -#include -#include -#else -#include -#include -#endif - -namespace at::native { -TORCH_IMPL_FUNC(lerp_Tensor_mps)(const Tensor& self, const Tensor& end, const Tensor& weight, const Tensor& out) { - TORCH_CHECK(out.is_mps()); - std::array args{{{out, "out", 0}, {self, "self", 1}, {end, "end", 2}, {weight, "weight", 3}}}; - checkAllSameGPU(__func__, args); - using namespace mps; - struct CachedGraph : public MPSCachedGraph { - CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {} - MPSGraphTensor* selfTensor_ = nil; - MPSGraphTensor* endTensor_ = nil; - MPSGraphTensor* weightTensor_ = nil; - MPSGraphTensor* outputTensor_ = nil; - }; - @autoreleasepool { - std::string key = "lerp_Tensor_mps" + getTensorsStringKey({self, end, weight}); - auto cachedGraph = LookUpOrCreateCachedGraph(key, [&](auto mpsGraph, auto graph) { - auto selfTensor = mpsGraphRankedPlaceHolder(mpsGraph, self); - auto endTensor = mpsGraphRankedPlaceHolder(mpsGraph, end); - auto weightTensor = mpsGraphRankedPlaceHolder(mpsGraph, weight); - auto distance = [mpsGraph subtractionWithPrimaryTensor:endTensor secondaryTensor:selfTensor name:nil]; - auto weighedDistance = [mpsGraph multiplicationWithPrimaryTensor:weightTensor secondaryTensor:distance name:nil]; - auto output = [mpsGraph additionWithPrimaryTensor:selfTensor secondaryTensor:weighedDistance name:nil]; - graph->selfTensor_ = selfTensor; - graph->endTensor_ = endTensor; - graph->weightTensor_ = weightTensor; - graph->outputTensor_ = output; - }); - auto selfPlaceholder = Placeholder(cachedGraph->selfTensor_, self); - auto endPlaceholder = Placeholder(cachedGraph->endTensor_, end); - auto weightPlaceholder = Placeholder(cachedGraph->weightTensor_, weight); - auto outputPlaceholder = Placeholder(cachedGraph->outputTensor_, out); - auto feeds = dictionaryFromPlaceholders(selfPlaceholder, endPlaceholder, weightPlaceholder); - runMPSGraph(getCurrentMPSStream(), cachedGraph->graph(), feeds, outputPlaceholder); - } -} - -} // namespace at::native diff --git a/aten/src/ATen/native/mps/operations/Linear.mm b/aten/src/ATen/native/mps/operations/Linear.mm index 219086edd8e37..2ccb89bc4d702 100644 --- a/aten/src/ATen/native/mps/operations/Linear.mm +++ b/aten/src/ATen/native/mps/operations/Linear.mm @@ -76,8 +76,6 @@ Tensor _mps_linear(const Tensor& input, const Tensor& weight_arg, const std::opt TORCH_CHECK(input.is_mps(), "Tensor for argument input is on ", input.device(), " but expected on mps"); TORCH_CHECK(supportedFloatingOrComplexType(weight_arg), "MPS device does not support linear for non-float weights"); TORCH_CHECK(weight_arg.is_mps(), "Tensor for argument weight is on ", weight_arg.device(), " but expected on mps"); - TORCH_CHECK((input.scalar_type() != kComplexFloat && input.scalar_type() != kComplexHalf), - "mps linear does not support complex types"); const Tensor& bias = *(at::borrow_from_optional_tensor(bias_opt)); const bool is_bias_defined = bias.defined(); @@ -117,8 +115,9 @@ Tensor _mps_linear(const Tensor& input, const Tensor& weight_arg, const std::opt // No-graph execution causes nonsense if these are non-contiguous. const bool is_contiguous = input.is_contiguous() && weight.is_contiguous() && bias.is_contiguous(); + const bool is_complex = input.is_complex() || weight.is_complex() || (is_bias_defined && bias.is_complex()); - if (is_macos_13_or_newer(MacOSVersion::MACOS_VER_15_0_PLUS) && is_contiguous) { + if (is_macos_13_or_newer(MacOSVersion::MACOS_VER_15_0_PLUS) && is_contiguous && !is_complex) { _mps_linear_nograph(input, weight, bias, output); // Squeeze last dim of 1D linear return weight_arg.dim() != 1 ? output : output.squeeze(-1); @@ -206,8 +205,7 @@ static Tensor _mps_linear_backward_input(IntArrayRef input_size, const Tensor& g MPSGraphTensor* outputTensor_ = nil; }; - Tensor output = at::empty( - input_size, grad_output.scalar_type(), std::nullopt, kMPS, std::nullopt, grad_output.suggest_memory_format()); + Tensor output = at::empty(input_size, grad_output.options()); TORCH_CHECK(output.is_mps()); if (grad_output.numel() == 0) { return output; @@ -275,18 +273,8 @@ static Tensor _mps_linear_backward_input(IntArrayRef input_size, const Tensor& g TORCH_CHECK(grad_output_reshaped.is_mps()); TORCH_CHECK(input_reshaped.is_mps()); - Tensor output = at::empty({grad_output_reshaped.size(1), input_reshaped.size(1)}, - grad_output.scalar_type(), - std::nullopt, - kMPS, - std::nullopt, - grad_output.suggest_memory_format()); - Tensor bias = at::empty({grad_output_reshaped.size(1)}, - grad_output.scalar_type(), - std::nullopt, - kMPS, - std::nullopt, - grad_output.suggest_memory_format()); + Tensor output = at::empty({grad_output_reshaped.size(1), input_reshaped.size(1)}, grad_output.options()); + Tensor bias = at::empty({grad_output_reshaped.size(1)}, grad_output.options()); TORCH_CHECK(output.is_mps()); TORCH_CHECK(bias.is_mps()); diff --git a/aten/src/ATen/native/mps/operations/LinearAlgebra.mm b/aten/src/ATen/native/mps/operations/LinearAlgebra.mm index ec48725891fb0..6bfd7d1049c4b 100644 --- a/aten/src/ATen/native/mps/operations/LinearAlgebra.mm +++ b/aten/src/ATen/native/mps/operations/LinearAlgebra.mm @@ -310,6 +310,20 @@ bool use_metal_mm(const Tensor& self, const Tensor& other, const Tensor& output) if (c10::isComplexType(self.scalar_type()) && self.size(1) > max_complex_inner_size) { return true; } + // Detect conditions that would trigger LORADOWN GEMV kernel with potential padding overflow + // See https://github.com/pytorch/pytorch/issues/178056 + if (self.scalar_type() == at::ScalarType::Half && (self.size(0) <= 16 || other.size(1) <= 16) && + self.stride(1) == 1 && other.stride(0) == 1) { + int64_t self_padding = self.stride(0) - self.size(1); + int64_t other_padding = other.stride(1) - other.size(0); + + if (self_padding > 15 || other_padding > 15 || self_padding % 4 != 0 || other_padding % 4 != 0) { + TORCH_WARN_ONCE( + "MPS mm implementation has a known issue with this shape, dtype and slice. Dispatching to metal implementation instead. This may impact performance."); + return true; + } + } + return !is_macos_14_4_or_newer && (self.stride(0) > max_stride_size || self.stride(1) > max_stride_size || self.size(0) > max_stride_size || self.size(1) > max_stride_size || other.stride(0) > max_stride_size || other.stride(1) > max_stride_size || @@ -698,13 +712,21 @@ static void linalg_inv_ex_out_mps_impl(const Tensor& A, bool check_errors, const }); // MPS TODO: // Strided API doesn't play nice with complex data types (at least not in case of matmul). + // MPSGraph's matrixMultiplication produces incorrect results with stride-0 NDArray + // inputs on macOS < 26.4 (only every 16th row is computed). Contiguify such tensors + // by disabling the strided API so they go through the gather/clone path first. + // See https://github.com/pytorch/pytorch/issues/180201 + static const bool is_macOS_26_4_or_newer = is_macos_13_or_newer(MacOSVersion::MACOS_VER_26_4_PLUS); + auto hasZeroStride = [](const Tensor& t) { + return std::ranges::any_of(t.strides(), [](auto s) { return s == 0; }); + }; + auto useStridedSelf = !isComplexType(self.scalar_type()) && (is_macOS_26_4_or_newer || !hasZeroStride(self)); + auto useStridedOther = !isComplexType(other.scalar_type()) && (is_macOS_26_4_or_newer || !hasZeroStride(other)); auto selfPlaceholder = self.numel() != 0 - ? Placeholder( - cachedGraph->inputTensor_, self, nil, true, MPSDataTypeInvalid, !isComplexType(self.scalar_type())) + ? Placeholder(cachedGraph->inputTensor_, self, nil, true, MPSDataTypeInvalid, useStridedSelf) : Placeholder(); auto otherPlaceholder = other.numel() != 0 - ? Placeholder( - cachedGraph->otherTensor_, other, nil, true, MPSDataTypeInvalid, !isComplexType(other.scalar_type())) + ? Placeholder(cachedGraph->otherTensor_, other, nil, true, MPSDataTypeInvalid, useStridedOther) : Placeholder(); auto outputPlaceholder = Placeholder(cachedGraph->outputTensor_, output); diff --git a/aten/src/ATen/native/mps/operations/Normalization.mm b/aten/src/ATen/native/mps/operations/Normalization.mm index 06cf207785b4d..d13ba8f5bf149 100644 --- a/aten/src/ATen/native/mps/operations/Normalization.mm +++ b/aten/src/ATen/native/mps/operations/Normalization.mm @@ -221,25 +221,29 @@ Check if running mean exists (maybe do this check before making graph) MPSGraphTensor* batchVarianceTensor = [mpsGraph varianceOfTensor:inputTensor axes:axes name:nil]; varTensor = batchVarianceTensor; if (has_running_mean) { + // Running stats may have a different dtype (e.g. float32 with float16 input) + auto running_mean_dtype = getMPSDataType(running_mean_opt.value()); // TODO: This is not the formula used in PyTorch, is this OK? Seems more robust // float besselCorrectionTerm = float(N) / std::max(N - 1.0f, 1.0f); float besselCorrectionTerm = float(N) / float(N - 1.0f); MPSGraphTensor* besselConstantTensor = [mpsGraph constantWithScalar:(double)besselCorrectionTerm shape:@[ @1 ] - dataType:input_mps_dtype]; - MPSGraphTensor* unbiasedVarianceTensor = [mpsGraph multiplicationWithPrimaryTensor:batchVarianceTensor - secondaryTensor:besselConstantTensor - name:nil]; + dataType:running_mean_dtype]; + MPSGraphTensor* unbiasedVarianceTensor = + [mpsGraph multiplicationWithPrimaryTensor:castMPSTensor(mpsGraph, batchVarianceTensor, running_mean_dtype) + secondaryTensor:besselConstantTensor + name:nil]; MPSGraphTensor* momentumTensor = [mpsGraph constantWithScalar:(double)momentum shape:@[ @1 ] - dataType:input_mps_dtype]; + dataType:running_mean_dtype]; MPSGraphTensor* oneMinusMomentum = [mpsGraph constantWithScalar:(double)(1.0 - momentum) shape:@[ @1 ] - dataType:input_mps_dtype]; + dataType:running_mean_dtype]; // Compute updated running mean - MPSGraphTensor* scaledBatchMean = [mpsGraph multiplicationWithPrimaryTensor:batchMeanTensor - secondaryTensor:momentumTensor - name:nil]; + MPSGraphTensor* scaledBatchMean = + [mpsGraph multiplicationWithPrimaryTensor:castMPSTensor(mpsGraph, batchMeanTensor, running_mean_dtype) + secondaryTensor:momentumTensor + name:nil]; MPSGraphTensor* scaledRunningMean = [mpsGraph multiplicationWithPrimaryTensor:runningMeanTensor secondaryTensor:oneMinusMomentum name:nil]; @@ -278,12 +282,16 @@ Check if running mean exists (maybe do this check before making graph) varTensor = saveVarTensor; } + // Cast weight and bias to input dtype if needed (mixed-precision support) + MPSGraphTensor* gammaTensor = has_weight ? castMPSTensor(mpsGraph, weightTensor, input_mps_dtype) : nil; + MPSGraphTensor* betaTensor = has_bias ? castMPSTensor(mpsGraph, biasTensor, input_mps_dtype) : nil; + // Compute output of batch norm MPSGraphTensor* outputTensor = [mpsGraph normalizationWithTensor:inputTensor meanTensor:saveMeanTensor varianceTensor:varTensor - gammaTensor:weightTensor - betaTensor:biasTensor + gammaTensor:gammaTensor + betaTensor:betaTensor epsilon:(float)epsilon name:nil]; @@ -610,7 +618,9 @@ Check if running mean exists (maybe do this check before making graph) NSString* ns_shape_key = [[input_shape valueForKey:@"description"] componentsJoinedByString:@","]; - std::string key = fmt::format("batch_norm_backward_mps:{}:{}:{}:{}:{}:{}:{}:{}", + auto input_mps_dtype = getMPSDataType(input); + auto weight_mps_dtype = has_weight ? getMPSDataType(weight_opt.value()) : input_mps_dtype; + std::string key = fmt::format("batch_norm_backward_mps:{}:{}:{}:{}:{}:{}:{}:{}:{}", get_mem_string(memory_format), epsilon, train, @@ -618,8 +628,8 @@ Check if running mean exists (maybe do this check before making graph) has_weight, [ns_shape_key UTF8String], c10::Join(",", grad_input_mask), - getMPSTypeString(input)); - auto input_mps_dtype = getMPSDataType(input); + getMPSTypeString(input), + has_weight ? getMPSTypeString(weight_opt.value()) : "none"); auto cachedGraph = LookUpOrCreateCachedGraph(key, [&](auto mpsGraph, auto newCachedGraph) { // NCHW - Channels dim is 1 int channelsDim = 1; @@ -628,8 +638,11 @@ Check if running mean exists (maybe do this check before making graph) // Shape is the ORIGINAL NCHW shape auto gradOutputTensor = mpsGraphRankedPlaceHolder(mpsGraph, getMPSDataType(grad_out), input_shape_readonly); MPSGraphTensor* weightTensor = nil; - if (has_weight) - weightTensor = mpsGraphRankedPlaceHolder(mpsGraph, getMPSDataType(weight_opt.value()), new_mean_shape); + MPSGraphTensor* weightTensorCasted = nil; + if (has_weight) { + weightTensor = mpsGraphRankedPlaceHolder(mpsGraph, weight_mps_dtype, new_mean_shape); + weightTensorCasted = castMPSTensor(mpsGraph, weightTensor, input_mps_dtype); + } MPSGraphTensor* runningMeanTensor = nil; MPSGraphTensor* runningVarTensor = nil; if (has_running_mean) { @@ -698,7 +711,7 @@ Check if running mean exists (maybe do this check before making graph) sourceTensor:inputTensor meanTensor:saveMeanTensor varianceTensor:revertSaveVarTensor - gammaTensor:weightTensor + gammaTensor:weightTensorCasted gammaGradientTensor:gradWeightTensor betaGradientTensor:gradBiasTensor reductionAxes:axes @@ -766,7 +779,7 @@ Check if running mean exists (maybe do this check before making graph) gradInputTensor = [mpsGraph multiplicationWithPrimaryTensor:unitTensor secondaryTensor:rsqrtTensor name:nil]; if (has_weight) gradInputTensor = [mpsGraph multiplicationWithPrimaryTensor:gradInputTensor - secondaryTensor:weightTensor + secondaryTensor:weightTensorCasted name:nil]; gradInputTensor = [mpsGraph multiplicationWithPrimaryTensor:gradInputTensor secondaryTensor:gradOutputTensor @@ -778,11 +791,13 @@ Check if running mean exists (maybe do this check before making graph) gradWeightTensor = [mpsGraph reshapeTensor:gradWeightTensor withShape:@[ input_shape_readonly[channelsDim] ] name:nil]; + gradWeightTensor = castMPSTensor(mpsGraph, gradWeightTensor, weight_mps_dtype); } if (grad_input_mask[2]) { gradBiasTensor = [mpsGraph reshapeTensor:gradBiasTensor withShape:@[ input_shape_readonly[channelsDim] ] name:nil]; + gradBiasTensor = castMPSTensor(mpsGraph, gradBiasTensor, weight_mps_dtype); } MPSGraphTensor* gradInputTensorFinal = nil; diff --git a/aten/src/ATen/native/mps/operations/RMSNorm.mm b/aten/src/ATen/native/mps/operations/RMSNorm.mm index 7948b5acd8e93..33348c6cbff8e 100644 --- a/aten/src/ATen/native/mps/operations/RMSNorm.mm +++ b/aten/src/ATen/native/mps/operations/RMSNorm.mm @@ -55,7 +55,7 @@ if (N <= LOOPED_LIMIT) { size_t threadgroup_needed = (N + N_READS - 1) / N_READS; size_t simds_needed = (threadgroup_needed + SIMD_SIZE - 1) / SIMD_SIZE; - size_t threadgroup_size = SIMD_SIZE * simds_needed; + threadgroup_size = SIMD_SIZE * simds_needed; assert(threadgroup_size <= maxThreadsPerGroup); } size_t n_threads = M * threadgroup_size; diff --git a/aten/src/ATen/native/mps/operations/Repeat.mm b/aten/src/ATen/native/mps/operations/Repeat.mm index f350b0137b05e..7abfb0a2e4500 100644 --- a/aten/src/ATen/native/mps/operations/Repeat.mm +++ b/aten/src/ATen/native/mps/operations/Repeat.mm @@ -10,6 +10,14 @@ #include #include +#ifndef AT_PER_OPERATOR_HEADERS +#include +#include +#else +#include +#include +#endif + namespace at::native { Tensor permute_mps(const Tensor& self, IntArrayRef dims) { @@ -36,7 +44,13 @@ Tensor repeat_mps(const Tensor& self, IntArrayRef repeats) { TORCH_CHECK(repeats.size() >= (size_t)self.dim(), "Number of dimensions of repeat dims can not be smaller than number of dimensions of tensor"); - TORCH_CHECK(!self.is_complex(), "repeat(): Not supported for complex yet!"); + + if (self.is_complex()) { + std::vector repeats_real = repeats.vec(); + repeats_real.push_back(1); + auto self_real = at::view_as_real(self); + return at::view_as_complex(repeat_mps(self_real, repeats_real)); + } // Add new leading dimensions to the tensor if the // number of target dimensions is larger than the diff --git a/aten/src/ATen/native/mps/operations/ScanKernel.h b/aten/src/ATen/native/mps/operations/ScanKernel.h new file mode 100644 index 0000000000000..06bedfbdd215e --- /dev/null +++ b/aten/src/ATen/native/mps/operations/ScanKernel.h @@ -0,0 +1,11 @@ +#pragma once + +namespace at::native::mps { + +void scan_simple_mps_impl( + const Tensor& self, + const Tensor& output, + int64_t dim, + const std::string& op_name); + +} // namespace at::native::mps diff --git a/aten/src/ATen/native/mps/operations/ScanKernel.mm b/aten/src/ATen/native/mps/operations/ScanKernel.mm index fce8f77e84f07..9b8a8a5d485f7 100644 --- a/aten/src/ATen/native/mps/operations/ScanKernel.mm +++ b/aten/src/ATen/native/mps/operations/ScanKernel.mm @@ -3,6 +3,7 @@ #include #include #include +#include #ifndef AT_PER_OPERATOR_HEADERS #include @@ -45,7 +46,7 @@ return {static_cast(grid_x), static_cast(grid_y)}; } -static void scan_simple_mps_impl(const Tensor& self, const Tensor& output, int64_t dim, const std::string& op_name) { +void scan_simple_mps_impl(const Tensor& self, const Tensor& output, int64_t dim, const std::string& op_name) { if (output.numel() == 0) { return; } diff --git a/aten/src/ATen/native/mps/operations/ScatterGather.mm b/aten/src/ATen/native/mps/operations/ScatterGather.mm index ce65421c71c9d..b9b055b4fee6f 100644 --- a/aten/src/ATen/native/mps/operations/ScatterGather.mm +++ b/aten/src/ATen/native/mps/operations/ScatterGather.mm @@ -10,10 +10,23 @@ #include #include #include +#include #endif namespace at::native { +static Tensor maybe_expand_0_dim(const Tensor& t) { + return t.dim() == 0 ? t.view({1}) : t; +} + +static Tensor expand_index_as_real(const Tensor& index) { + auto index_view = maybe_expand_0_dim(index); + std::vector index_expanded_sizes = index_view.sizes().vec(); + index_expanded_sizes.push_back(2); + auto index_expanded = index_view.unsqueeze(-1).expand(index_expanded_sizes); + return index_expanded; +} + TORCH_IMPL_FUNC(gather_out_mps) (const Tensor& self_arg, int64_t dim, const Tensor& index, bool sparse_grad, const Tensor& output) { using namespace mps; @@ -27,7 +40,14 @@ TORCH_CHECK(!sparse_grad, "sparse_grad not supported in MPS yet") TORCH_CHECK(self.scalar_type() == output.scalar_type(), "gather(): self and output must have the same scalar type"); TORCH_CHECK(dim >= 0 && dim < self.dim(), "gather(): Indexing dim ", dim, " is out of bounds of tensor"); - TORCH_CHECK(!self.is_complex(), "gather(): Yet not supported for complex"); + + if (self.is_complex()) { + auto self_real = at::view_as_real(self); + auto index_expanded = expand_index_as_real(index); + auto output_real = at::view_as_real(maybe_expand_0_dim(output)); + structured_gather_out_mps::impl(self_real, dim, index_expanded, sparse_grad, output_real); + return; + } struct CachedGraph : public MPSCachedGraph { CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {} @@ -145,7 +165,15 @@ static void scatter_mps_general(const Tensor& self_arg, TORCH_CHECK(self.scalar_type() == output.scalar_type() && output.scalar_type() == src.scalar_type(), "scatter(): self, src and output must have the same scalar type"); TORCH_CHECK(dim >= 0 && dim < self.dim(), "scatter(): Indexing dim ", dim, " is out of bounds of tensor"); - TORCH_CHECK(!self.is_complex(), "scatter(): Yet not supported for complex"); + + if (self.is_complex()) { + auto self_real = at::view_as_real(self); + auto index_expanded = expand_index_as_real(index); + auto src_real = at::view_as_real(maybe_expand_0_dim(src)); + auto output_real = at::view_as_real(maybe_expand_0_dim(output)); + scatter_mps_general(self_real, dim, index_expanded, src_real, output_real, func_name, reduce); + return; + } struct CachedGraph : public MPSCachedGraph { CachedGraph(MPSGraph* graph) : MPSCachedGraph(graph) {} diff --git a/aten/src/ATen/native/mps/operations/UnaryOps.mm b/aten/src/ATen/native/mps/operations/UnaryOps.mm index eb827d9f44c26..9dc318a393a5c 100644 --- a/aten/src/ATen/native/mps/operations/UnaryOps.mm +++ b/aten/src/ATen/native/mps/operations/UnaryOps.mm @@ -3,6 +3,7 @@ #include #include #include +#include #ifndef AT_PER_OPERATOR_HEADERS #include @@ -333,8 +334,21 @@ static void cumulative_op_impl(const Tensor& self, "(original dim is ", dim, ")"); - TORCH_CHECK(!self.is_complex(), "cumulative ops are not yet supported for complex"); auto input = dtype.has_value() ? self.to(dtype.value()) : self; + if (input.is_complex()) { + if (cumulativeOpType == MPSCumulativeOpType::CUMSUM) { + auto input_real = at::view_as_real(input.dim() == 0 ? input.view({1}) : input); + auto result_real = at::view_as_real(result.dim() == 0 ? result.view({1}) : result); + return cumulative_op_impl( + input_real, wrapped_dim, std::nullopt, result_real, MPSCumulativeOpType::CUMSUM, "cumsum_out_mps"); + } else if (cumulativeOpType == MPSCumulativeOpType::CUMPROD) { + auto input_view = input.dim() == 0 ? input.view({1}) : input; + auto result_view = result.dim() == 0 ? result.view({1}) : result; + return mps::scan_simple_mps_impl(input_view, result_view, wrapped_dim, "cumprod"); + } else { + TORCH_INTERNAL_ASSERT(false); + } + } // issue #103810551: cumsum / cumprod are broken for int8, int16 and as chances for overflow are pretty high, cast to // int32 fixed in macOS 13.3 diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index ec537575032ae..9e39eda0368fe 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -2736,8 +2736,7 @@ device_check: NoCheck # TensorIterator variants: function, method dispatch: - CPU, CUDA: fill_ - MPS: fill_scalar_mps + CPU, CUDA, MPS: fill_ QuantizedCPU, QuantizedCUDA: fill_quantized_ Meta: fill_meta_ SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: fill_sparse_csr_ @@ -2748,8 +2747,7 @@ device_check: NoCheck # TensorIterator variants: function, method dispatch: - CPU, CUDA: fill_ - MPS: fill_tensor_mps_ + CPU, CUDA, MPS: fill_ QuantizedCPU, QuantizedCUDA: fill_quantized_ Meta: fill_meta_ NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: fill_nested_ @@ -2949,6 +2947,8 @@ autogen: _grid_sampler_2d_cpu_fallback.out - func: _grid_sampler_2d_cpu_fallback_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> (Tensor, Tensor) + dispatch: + CompositeExplicitAutograd: _grid_sampler_2d_cpu_fallback_backward - func: grid_sampler_3d(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor dispatch: @@ -3297,6 +3297,42 @@ device_guard: False manual_cpp_binding: True +- func: numel(Tensor self) -> int + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: dim(Tensor self) -> int + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: get_device(Tensor self) -> int + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: storage_offset(Tensor self) -> int + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: is_contiguous(Tensor self) -> bool + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: is_contiguous.memory_format(Tensor self, MemoryFormat memory_format) -> bool + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + - func: kl_div(Tensor self, Tensor target, int reduction=Mean, *, bool log_target=False) -> Tensor - func: kron(Tensor self, Tensor other) -> Tensor @@ -3888,7 +3924,7 @@ device_check: NoCheck # TensorIterator structured: True dispatch: - CPU, CUDA, MTIA: aminmax_out + CPU, CUDA: aminmax_out MPS: aminmax_out_mps tags: reduction @@ -3943,7 +3979,7 @@ - func: amax.out(Tensor self, int[1] dim=[], bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) structured: True dispatch: - CPU, CUDA, MTIA: amax_out + CPU, CUDA: amax_out MPS: amax_out_mps tags: reduction @@ -4132,7 +4168,7 @@ - func: amin.out(Tensor self, int[1] dim=[], bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) structured: True dispatch: - CPU, CUDA, MTIA: amin_out + CPU, CUDA: amin_out MPS: amin_out_mps tags: reduction @@ -4421,7 +4457,7 @@ - func: mvlgamma.out(Tensor self, int p, *, Tensor(a!) out) -> Tensor(a!) dispatch: - CPU, CUDA: mvlgamma_out + CPU, CUDA, MPS: mvlgamma_out tags: pointwise - func: mvlgamma(Tensor self, int p) -> Tensor @@ -5417,8 +5453,7 @@ structured_inherits: TensorIteratorBase python_module: nn dispatch: - CPU, CUDA, MTIA: silu_out - MPS: silu_out_mps + CPU, CUDA, MPS, MTIA: silu_out tags: pointwise - func: silu_backward.grad_input(Tensor grad_output, Tensor self, *, Tensor(a!) grad_input) -> Tensor(a!) @@ -5426,8 +5461,7 @@ structured_inherits: TensorIteratorBase python_module: nn dispatch: - CPU, CUDA: silu_backward_out - MPS: silu_backward_out_mps + CPU, CUDA, MPS: silu_backward_out tags: pointwise - func: silu_backward(Tensor grad_output, Tensor self) -> Tensor @@ -6660,7 +6694,6 @@ dispatch: CPU, CUDA: var MPS: var_mps - MTIA: var_mtia tags: [core, reduction] - func: var.out(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) @@ -6822,6 +6855,7 @@ dispatch: CPU: _standard_gamma_grad_cpu CUDA: _standard_gamma_grad_cuda + MPS: _standard_gamma_grad_mps autogen: _standard_gamma_grad.out - func: _standard_gamma(Tensor self, Generator? generator=None) -> Tensor @@ -6829,9 +6863,32 @@ dispatch: CPU: _s_gamma_cpu CUDA: _s_gamma_cuda + MPS: _s_gamma_mps tags: nondeterministic_seeded autogen: _standard_gamma.out +- func: _philox_key_split(Tensor key, int num_splits) -> Tensor + variants: function + dispatch: + CUDA: _philox_key_split_cuda + +- func: _philox_key_fold_in(Tensor key, int data) -> Tensor + variants: function + dispatch: + CUDA: _philox_key_fold_in_cuda + +- func: _philox_normal_(Tensor(a!) self, Tensor key, float mean=0, float std=1) -> Tensor(a!) + variants: function, method + dispatch: + CUDA: _philox_normal_cuda_ + autogen: _philox_normal, _philox_normal.out + +- func: _philox_uniform_(Tensor(a!) self, Tensor key, float low=0, float high=1) -> Tensor(a!) + variants: function, method + dispatch: + CUDA: _philox_uniform_cuda_ + autogen: _philox_uniform, _philox_uniform.out + - func: _dirichlet_grad(Tensor x, Tensor alpha, Tensor total) -> Tensor dispatch: CPU: _dirichlet_grad_cpu @@ -7120,8 +7177,7 @@ device_check: NoCheck # TensorIterator variants: method, function dispatch: - CPU, CUDA: zero_ - MPS: zero_mps_ + CPU, CUDA, MPS: zero_ Meta: zero_meta_ SparseCPU, SparseCUDA, SparseMPS, SparseMeta: zero_sparse_ SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: zero_sparse_csr_ @@ -7329,6 +7385,7 @@ - func: _scaled_mm_v2(Tensor self, Tensor mat2, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? bias, ScalarType? out_dtype, int[] contraction_dim=[], bool use_fast_accum=False) -> Tensor variants: function dispatch: + CPU: _scaled_mm_cpu_v2 CUDA: _scaled_mm_cuda_v2 XPU: _scaled_mm_xpu_v2 tags: needs_exact_strides @@ -7336,6 +7393,7 @@ - func: _scaled_mm_v2.out(Tensor self, Tensor mat2, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? bias, ScalarType? out_dtype, int[] contraction_dim=[], bool use_fast_accum=False, *, Tensor(a!) out) -> Tensor(a!) variants: function dispatch: + CPU: _scaled_mm_cpu_v2_out CUDA: _scaled_mm_cuda_v2_out XPU: _scaled_mm_xpu_v2_out tags: needs_exact_strides @@ -9570,12 +9628,14 @@ dispatch: CPU: nonzero_static_out_cpu CUDA: nonzero_static_out_cuda + MPS: nonzero_static_out_mps - func: nonzero_static(Tensor self, *, SymInt size, int fill_value=-1) -> Tensor variants: method, function dispatch: CPU: nonzero_static_cpu CUDA: nonzero_static_cuda + MPS: nonzero_static_mps - func: nonzero_numpy(Tensor self) -> Tensor[] variants: method, function @@ -10012,8 +10072,7 @@ structured: True structured_inherits: TensorIteratorBase dispatch: - CPU, CUDA: lerp_Tensor - MPS: lerp_Tensor_mps + CPU, CUDA, MPS: lerp_Tensor tags: pointwise - func: lerp.Scalar(Tensor self, Tensor end, Scalar weight) -> Tensor @@ -11798,6 +11857,14 @@ CUDA: foreach_tensor_zero_cuda_ autogen: _foreach_zero, _foreach_zero.out +- func: _foreach_clone(Tensor[] self, *, MemoryFormat? memory_format=None) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clone_slow + CUDA: foreach_tensor_clone_cuda + autogen: _foreach_clone.out + - func: _foreach_copy_(Tensor(a!)[] self, Tensor[] src, bool non_blocking=False) -> () device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices variants: function @@ -12956,6 +13023,10 @@ python_module: nn autogen: _upsample_bicubic2d_aa.vec_out +- func: _upsample_lanczos2d_aa.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor + python_module: nn + autogen: _upsample_lanczos2d_aa.vec_out + - func: upsample_nearest1d.vec(Tensor input, SymInt[]? output_size, float[]? scale_factors) -> Tensor python_module: nn autogen: upsample_nearest1d.vec_out @@ -13102,6 +13173,26 @@ python_module: nn structured_delegate: _upsample_bicubic2d_aa_backward.grad_input +- func: _upsample_lanczos2d_aa.out(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_lanczos2d_aa_out_cpu + +- func: _upsample_lanczos2d_aa(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_lanczos2d_aa.out + +- func: _upsample_lanczos2d_aa_backward.grad_input(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_lanczos2d_aa_backward_out_cpu + +- func: _upsample_lanczos2d_aa_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_lanczos2d_aa_backward.grad_input + - func: upsample_trilinear3d.out(Tensor self, SymInt[3] output_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) python_module: nn structured: True @@ -15185,7 +15276,7 @@ variants: function tags: nondeterministic_seeded -- func: _scaled_dot_product_attention_math_for_mps(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, Tensor? dropout_mask=None, *, float? scale=None) -> (Tensor, Tensor) +- func: _scaled_dot_product_attention_math_for_mps(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, Tensor? dropout_mask=None, *, float? scale=None, bool enable_gqa=False) -> (Tensor, Tensor) dispatch: MPS: _scaled_dot_product_attention_math_mps tags: nondeterministic_seeded diff --git a/aten/src/ATen/native/sparse/SparseCsrTensor.cpp b/aten/src/ATen/native/sparse/SparseCsrTensor.cpp index 5ecfbd589ca39..79ce7fb17bfb3 100644 --- a/aten/src/ATen/native/sparse/SparseCsrTensor.cpp +++ b/aten/src/ATen/native/sparse/SparseCsrTensor.cpp @@ -189,8 +189,10 @@ static void _validate_sparse_compressed_tensor_args_worker(const Tensor& compres batch_ndim, " + ", block_ndim, ") but got ", values.dim()); // 3.5 - TORCH_CHECK(plain_indices.stride(-1) == 1, - "expected ", plain_indices_name, " to be a contiguous tensor per batch"); + if (plain_indices.numel() != 0) { + TORCH_CHECK(plain_indices.stride(-1) == 1, + "expected ", plain_indices_name, " to be a contiguous tensor per batch"); + } // 3.6 TORCH_CHECK(compressed_indices.stride(-1) == 1, diff --git a/aten/src/ATen/native/sparse/cuda/SparseCUDAApplyUtils.cuh b/aten/src/ATen/native/sparse/cuda/SparseCUDAApplyUtils.cuh index 078b354659d35..7b9e61ceb4cc7 100644 --- a/aten/src/ATen/native/sparse/cuda/SparseCUDAApplyUtils.cuh +++ b/aten/src/ATen/native/sparse/cuda/SparseCUDAApplyUtils.cuh @@ -246,11 +246,7 @@ __global__ void coalesceValuesKernel( // `if constexpr` when CUDA codes will be compiled under C++-17, see // gh-56055 for blockers. template -#ifdef USE_ROCM -C10_LAUNCH_BOUNDS_1(C10_WARP_SIZE_STATIC*4) -#else -C10_LAUNCH_BOUNDS_1(C10_WARP_SIZE*4) -#endif +C10_LAUNCH_BOUNDS_1(C10_WARP_SIZE_UPPER_BOUND*4) __global__ void coalesceValuesKernel( int64_t *segment_offsets, int64_t *value_indices, bool *values, bool *newValues, diff --git a/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp b/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp index 9d735ac0f2c88..24893da98c6c3 100644 --- a/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp +++ b/aten/src/ATen/native/sparse/cuda/cuSPARSELtOps.cpp @@ -84,7 +84,7 @@ at::Tensor _cslt_compress(const Tensor& sparse_input) { type = CUDA_R_32F; break; #endif -#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 && !defined(USE_ROCM) +#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 || defined(USE_ROCM) case at::ScalarType::Float8_e4m3fn: type = CUDA_R_8F_E4M3; break; @@ -203,12 +203,18 @@ std::tuple _cslt_sparse_mm_impl( compute_type = CUSPARSE_COMPUTE_32F; break; #endif -// if cuSPARSELt >= 6.2.3, we can add Float8 support -#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 && !defined(USE_ROCM) +// cuSPARSELt >= 0.6.2 or hipSparseLt: add Float8 support +#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 || defined(USE_ROCM) case at::ScalarType::Float8_e4m3fn: input_type = CUDA_R_8F_E4M3; +#ifdef USE_ROCM + // hipSparseLt 0.2.7: FP8 input only supports FP32 output + output_type = CUDA_R_32F; + C_type = CUDA_R_32F; +#else output_type = CUDA_R_8F_E4M3; C_type = CUDA_R_16F; +#endif compute_type = CUSPARSE_COMPUTE_32F; break; #endif @@ -265,10 +271,11 @@ std::tuple _cslt_sparse_mm_impl( break; } } -// cslt 0.6.2+: fp8 fp8 -> {fp8, fp16, bf16, fp32} support -#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 && !defined(USE_ROCM) +// cslt 0.6.2+ or hipSparseLt: fp8 output dtype support +#if defined(CUSPARSELT_VERSION) && CUSPARSELT_VERSION >= 602 || defined(USE_ROCM) else if (input_type == CUDA_R_8F_E4M3) { switch (out_dtype) { +#ifndef USE_ROCM case at::ScalarType::Float8_e4m3fn: output_type = CUDA_R_8F_E4M3; C_type = CUDA_R_16F; @@ -281,6 +288,7 @@ std::tuple _cslt_sparse_mm_impl( output_type = CUDA_R_16BF; C_type = CUDA_R_16BF; break; +#endif case at::ScalarType::Float: output_type = CUDA_R_32F; C_type = CUDA_R_32F; @@ -288,7 +296,11 @@ std::tuple _cslt_sparse_mm_impl( default: TORCH_CHECK( false, - "Unsupported out_dtype passed, must be one of {fp16, bf16, float32} for fp8 inputs"); +#ifdef USE_ROCM + "Unsupported out_dtype passed, must be float32 for fp8 inputs on ROCm"); +#else + "Unsupported out_dtype passed, must be one of {fp8, fp16, bf16, float32} for fp8 inputs"); +#endif break; } } @@ -468,6 +480,8 @@ std::tuple _cslt_sparse_mm_impl( &alg_id, sizeof(alg_id))); +#ifndef USE_ROCM + // hipSPARSELt does not support querying SPLIT_K attributes TORCH_CUDASPARSE_CHECK(cusparseLtMatmulAlgGetAttribute( &handle, &alg_sel, @@ -481,6 +495,7 @@ std::tuple _cslt_sparse_mm_impl( CUSPARSELT_MATMUL_SPLIT_K_MODE, &splitKMode, sizeof(splitKMode))); +#endif TORCH_CUDASPARSE_CHECK(cusparseLtMatmulAlgGetAttribute( &handle, diff --git a/aten/src/ATen/native/transformers/attention.cpp b/aten/src/ATen/native/transformers/attention.cpp index af8487bb05287..eecfe1058c2c4 100644 --- a/aten/src/ATen/native/transformers/attention.cpp +++ b/aten/src/ATen/native/transformers/attention.cpp @@ -492,6 +492,20 @@ int64_t _fused_sdp_choice_meta( return choice_int; } #endif + bool has_xpu = query_key_set.has(c10::DispatchKey::XPU); + if (has_xpu) { + auto choice_int = _fused_sdp_choice_stub( + at::kXPU, + query_, + key, + value, + attn_mask_, + dropout_p, + is_causal, + scale, + enable_gqa); + return choice_int; + } return static_cast(sdp::SDPBackend::math); } namespace { @@ -676,11 +690,11 @@ Tensor _safe_softmax( // greater than 0.0 is specified. // // Args: -// query (Tensor): Query tensor; shape (N, ..., L, E) -// key (Tensor): Key tensor; shape (N, ..., S, E) -// value (Tensor): Value tensor; shape (N, ..., S, E) +// query (Tensor): Query tensor; shape (N, ..., Hq, L, E) +// key (Tensor): Key tensor; shape (N, ..., H, S, E) +// value (Tensor): Value tensor; shape (N, ..., H, S, Ev) // attn_mask (optional Tensor): Attention mask; shape must be broadcastable to the shape of attention weights, -// which is (N,..., L, S). Two types of masks are supported. +// which is (N,..., Hq, L, S). Two types of masks are supported. // A boolean mask where a value of True indicates that the element *should* take part in attention. // A float mask of the same type as query, key, value that is added to the attention score. // dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied @@ -692,14 +706,17 @@ Tensor _safe_softmax( // sparse masks) via tensor subclassing, allowing for a leaner API. // // Returns a tensor: -// output (Tensor): Attention output; shape (N, ..., L, E) +// output (Tensor): Attention output; shape (N, ..., Hq, L, Ev) // // Shape legend: // N: Batch size // ...: Any number of other batch dimensions (optional) // S: Source sequence length // L: Target sequence length -// E: Embedding dimension +// E: Embedding dimension of the query and key +// Ev: Embedding dimension of the value +// Hq: Number of heads of query +// H: Number of heads of key and value Tensor scaled_dot_product_attention( const Tensor& query_, const Tensor& key, @@ -775,48 +792,8 @@ Tensor scaled_dot_product_attention( return std::get<0>(out_lse_softmax); } case SDPBackend::math: { -#ifdef USE_MPS - TORCH_CHECK_NOT_IMPLEMENTED( - c10::isFloatingType(query_.scalar_type()), - "scaled_dot_product_attention for MPS does not support dtype ", - query_.scalar_type()); - TORCH_CHECK_NOT_IMPLEMENTED( - c10::isFloatingType(key.scalar_type()), - "scaled_dot_product_attention for MPS does not support dtype ", - key.scalar_type()); - TORCH_CHECK_NOT_IMPLEMENTED( - c10::isFloatingType(value.scalar_type()), - "scaled_dot_product_attention for MPS does not support dtype ", - value.scalar_type()); - const auto any_nested = query_.is_nested() || key.is_nested() || value.is_nested(); const bool any_inputs_require_grad = query_.requires_grad() || key.requires_grad() || value.requires_grad(); - const auto all_contiguous = query_.is_contiguous_or_false() && key.is_contiguous_or_false() && value.is_contiguous_or_false(); - if (query_device_type == DeviceType::MPS && dropout_p == 0.0 - && !(GradMode::is_enabled() && any_inputs_require_grad) - && (all_contiguous || mps::is_macos_13_or_newer(mps::MacOSVersion::MACOS_VER_15_0_PLUS)) - && !any_nested) { - if (enable_gqa) { - int64_t q_heads = query_.size(-3); - int64_t k_heads = key.size(-3); - int64_t repeat_factor = q_heads / k_heads; - - if (repeat_factor > 1) { - TORCH_CHECK(q_heads % k_heads == 0, - "For GQA, the query tensor's head dimension (" + std::to_string(q_heads) + - ") must be divisible by the key tensor's head dimension (" + std::to_string(k_heads) + ")."); - auto repeated_key = key.repeat_interleave(repeat_factor, /*dim=*/-3); - auto repeated_value = value.repeat_interleave(repeat_factor, /*dim=*/-3); - return std::get<0>(at::_scaled_dot_product_attention_math_for_mps( - query_, - repeated_key, - repeated_value, - attn_mask, - dropout_p, - is_causal, - std::nullopt, /*dropout_mask*/ - scale)); - } - } + if (query_device_type == c10::kMPS && !(at::GradMode::is_enabled() && any_inputs_require_grad)) { return std::get<0>(at::_scaled_dot_product_attention_math_for_mps( query_, key, @@ -825,9 +802,9 @@ Tensor scaled_dot_product_attention( dropout_p, is_causal, std::nullopt, /*dropout_mask*/ - scale)); + scale, + enable_gqa)); } -#endif return std::get<0>(at::_scaled_dot_product_attention_math( query_, key, diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip b/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip index b96e80d5e5a9e..f3816296a1d4f 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/aot/mha_all_aot.hip @@ -318,8 +318,16 @@ mha_fwd_aot(const at::Tensor &q, // batch_size x seqlen_q x num_heads x persistent_counter, stream); } - - return {out, q_padded, k_padded, v_padded, M.view({batch_size, num_heads, seqlen_q}), seed_t, offset_t, softmax_fa_t}; + // Note: These are propagated up to the return of mha_fwd(). comments + // represent the assignments at that level + return {out, // output + q_padded, // q_padded + k_padded, // k_padded + v_padded, // v_padded + M.view({batch_size, num_heads, seqlen_q}), // logsumexp + seed_t, // philox_seed + offset_t, // philox_offset + softmax_fa_t};// debug_attn_mask } std::tuple @@ -635,7 +643,7 @@ mha_bwd_aot(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x hea TORCH_CHECK(dv.stride(-1) == 1, "dv must have contiguous last dimension"); CHECK_SHAPE(dv, batch_size, seqlen_k, num_heads_k, head_size); } else { - dv = at::empty_like(k); + dv = at::empty_like(v); } auto [needs_swa, window_left, window_right] = calculate_swa(window_size_left, diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip index 2d3692d9f98df..d7633521d8eef 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_bwd_ck.hip @@ -127,7 +127,7 @@ aiter::mha_bwd_args get_ck_fmha_bwd_args(const mask_info &mask, return aiter::mha_bwd_args{ // aiter args static_cast(mask.type), - true, // use_asm_v3 + hdim <= 192, // use_asm_v3: ASM v3 only supports head dim <= 192 true, // v3_atomic_fp32 1, // v3_bf16_cvt false, // v3_api_check @@ -375,11 +375,16 @@ mha_bwd_ck(const at::Tensor &dout, // batch_size x seqlen_q x dv_expanded = dv; } - uint64_t drop_seed = 1, drop_offset = 0; - drop_seed = *philox_seed.data_ptr(); - drop_offset = *philox_offset.data_ptr(); - auto drop_seed_offset = std::make_pair(&drop_seed, &drop_offset); + auto gen = at::get_generator_or_default( + std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); + uint64_t* drop_seed, drop_offset; + int64_t counter_offset = batch_size * num_heads * ck_tile::get_warp_size(); + std::pair drop_seed_offset = {nullptr,nullptr}; + if(is_dropout) { + drop_seed_offset.first = philox_seed.data_ptr(); + drop_seed_offset.second = philox_offset.data_ptr(); + } if (seqlen_q > 0) { ck_tile::stream_config stream_config{stream}; diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip index 441589e70d763..c5902e33afd36 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_fwd_ck.hip @@ -183,7 +183,6 @@ mha_fwd_ck(const at::Tensor &q, // batch_size x seqlen_q x TORCH_CHECK(v.stride(-1) == 1, "Input tensor must have contiguous last dimension"); const auto sizes = q.sizes(); - const int batch_size = sizes[0]; int seqlen_q = sizes[1]; int num_heads = sizes[2]; @@ -232,7 +231,6 @@ mha_fwd_ck(const at::Tensor &q, // batch_size x seqlen_q x CHECK_SHAPE(k, batch_size, seqlen_k, num_heads_k, head_size); CHECK_SHAPE(v, batch_size, seqlen_k, num_heads_k, head_size); - at::Tensor q_padded, k_padded, v_padded; if (head_size % 8 != 0) { q_padded = at::pad(temp_q, {0, 8 - head_size % 8}); @@ -245,7 +243,6 @@ mha_fwd_ck(const at::Tensor &q, // batch_size x seqlen_q x v_padded = v; } - at::Tensor out; if (out_.has_value()) { out = out_.value(); @@ -272,7 +269,6 @@ mha_fwd_ck(const at::Tensor &q, // batch_size x seqlen_q x auto opts = q.options(); bool has_lse = true; bool has_dropout = p_dropout > 0.0f; - at::Tensor softmax_lse; // TODO - check gradient, only training require lse softmax_lse = at::empty({batch_size, num_heads, seqlen_q}, opts.dtype(at::kFloat)); @@ -283,47 +279,46 @@ mha_fwd_ck(const at::Tensor &q, // batch_size x seqlen_q x p = at::empty({batch_size, num_heads, seqlen_q, seqlen_k}, opts.dtype(at::kByte)); } else { - p = at::empty({ 0 }, opts); + p = at::empty({ 0, 0, 0, 0 }, opts.dtype(at::kByte)); } + uint64_t drop_seed = 1, drop_offset = 0; + at::Tensor seed_t, offset_t; int64_t counter_offset = batch_size * num_heads * ck_tile::get_warp_size(); - auto rng_state = at::empty({2}, opts.dtype(at::kLong)); - auto rng_state_ptr = reinterpret_cast(rng_state.data_ptr()); - - at::Tensor seed_t, offset_t; + // rng_state is used to pass philox params to CK in a type it likes i.e. uint64 + auto rng_state_options = at::TensorOptions().dtype(at::kUInt64).device(at::kCUDA); + auto rng_state = at::zeros({2}, rng_state_options.dtype(at::kUInt64)); + auto _unused = at::empty({}, at::dtype(c10::kUInt64).device(at::kCUDA)); if (p_dropout > 0.0) { + auto gen = at::get_generator_or_default( gen_, at::cuda::detail::getDefaultCUDAGenerator()); + // See Note [Acquire lock when using random generators] std::lock_guard lock(gen->mutex_); - auto philox_args = gen->philox_cuda_state(counter_offset); + std::tie(drop_seed, drop_offset) = at::cuda::philox::unpack(philox_args); + rng_state[0] = *(reinterpret_cast(&drop_seed)); + rng_state[1] = *(reinterpret_cast(&drop_offset)); - - hipLaunchKernelGGL( - flash::ParsePhiloxCudaState, dim3(1), dim3(64), 0, at::cuda::getCurrentCUDAStream(), philox_args, rng_state_ptr); - seed_t = at::scalar_tensor(at::Scalar(static_cast(rng_state_ptr[0])), at::dtype(at::kLong)); - offset_t = at::scalar_tensor(at::Scalar(static_cast(rng_state_ptr[1])), at::dtype(at::kLong)); - } - else - { + } else { seed_t = at::empty({}, at::dtype(at::kLong).device(at::kCUDA)); offset_t = at::empty({}, at::dtype(at::kLong).device(at::kCUDA)); } - std::optional attn_bias; + std::pair drop_seed_offset; if( attn_bias_.has_value()) { attn_bias = attn_bias_; } - if (seqlen_k > 0) { - auto drop_seed_offset = std::make_pair(rng_state_ptr, rng_state_ptr + 1); - auto stream = at::cuda::getCurrentCUDAStream().stream(); + drop_seed_offset.first = rng_state[0].data_ptr(); + drop_seed_offset.second = rng_state[1].data_ptr(); + auto stream = at::cuda::getCurrentHIPStream().stream(); ck_tile::stream_config stream_config{stream}; auto traits = @@ -364,12 +359,28 @@ mha_fwd_ck(const at::Tensor &q, // batch_size x seqlen_q x out.zero_(); softmax_lse.fill_(std::numeric_limits::infinity()); } - if (seqlenq_ngroups_swapped) { out = out.transpose(1, 2).reshape({batch_size, 1, num_heads_k * seqlen_q, head_size}); q_padded = q_padded.transpose(1, 2).reshape({batch_size, 1, num_heads_k * seqlen_q, head_size}); softmax_lse = softmax_lse.reshape({batch_size, num_heads_k * seqlen_q, 1}); } - return {out, q_padded, k_padded, v_padded, softmax_lse, seed_t, offset_t, p}; + // Before returning populate seed_t and offset_t if dropout is used + if(p_dropout > 0.0) + { + seed_t = at::scalar_tensor(at::Scalar(reinterpret_cast(*(drop_seed_offset.first))), + at::dtype(at::kUInt64).device(at::kCUDA)); + offset_t = at::scalar_tensor(at::Scalar(reinterpret_cast(*(drop_seed_offset.second))), + at::dtype(at::kUInt64).device(at::kCUDA)); + } + // Note: These are propagated up to the return of mha_fwd(). comments + // represent the assignments at that level + return {out, // output + q_padded, // q_padded + k_padded, // k_padded + v_padded, // v_padded + softmax_lse, // logsumexp + seed_t, // philox_seed + offset_t, // philox_offset + p}; // debug_attn_mask } } //namespace pytorch_flash diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_bwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_bwd_ck.hip index 3031965833608..896afc7320e96 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_bwd_ck.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_bwd_ck.hip @@ -401,8 +401,10 @@ mha_varlen_bwd_ck(const at::Tensor &dout, // total_q x num_hea uint64_t drop_seed = 1, drop_offset = 0; - drop_seed = *philox_seed.data_ptr(); - drop_offset = *philox_offset.data_ptr(); + if (is_dropout) { + drop_seed = *reinterpret_cast(philox_seed.data_ptr()); + drop_offset = *reinterpret_cast(philox_offset.data_ptr()); + } auto drop_seed_offset = std::make_pair(&drop_seed, &drop_offset); if (max_seqlen_q > 0) { diff --git a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip index 2a9d4899e8236..39ce82ef1e700 100644 --- a/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip +++ b/aten/src/ATen/native/transformers/hip/flash_attn/ck/mha_varlen_fwd_ck.hip @@ -73,7 +73,7 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, ck_tile::index_t nhead_stride_k = k.stride(1); ck_tile::index_t nhead_stride_v = v.stride(1); ck_tile::index_t nhead_stride_o = out.stride(1); - ck_tile::index_t nhead_stride_lse = has_lse ? softmax_lse.stride(1) : 0; + ck_tile::index_t nhead_stride_lse = has_lse ? softmax_lse.stride(0) : 0; ck_tile::index_t nhead_stride_randval = has_dropout_randval ? dropout_randval.stride(0) : 0; ck_tile::index_t batch_stride_q = 0; @@ -81,7 +81,7 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse, ck_tile::index_t batch_stride_v = 0; ck_tile::index_t batch_stride_o = 0; - ck_tile::index_t batch_stride_lse = has_lse ? softmax_lse.stride(0) : 0; + ck_tile::index_t batch_stride_lse = 0; ck_tile::index_t batch_stride_randval = 0; void *attn_bias_ptr = nullptr; @@ -287,7 +287,7 @@ mha_varlen_fwd_ck(const at::Tensor &q, // total_q x num_heads at::Tensor softmax_lse; // TODO - check gradient, only training require lse - softmax_lse = at::empty({batch_size, num_heads, max_seqlen_q}, opts.dtype(at::kFloat)); + softmax_lse = at::empty({num_heads, total_q}, opts.dtype(at::kFloat)); at::Tensor p; if (return_dropout_randval) { diff --git a/aten/src/ATen/native/transformers/xpu/sdp_utils.cpp b/aten/src/ATen/native/transformers/xpu/sdp_utils.cpp index a69ecf3274d2a..b5abfa27e8baa 100644 --- a/aten/src/ATen/native/transformers/xpu/sdp_utils.cpp +++ b/aten/src/ATen/native/transformers/xpu/sdp_utils.cpp @@ -29,7 +29,8 @@ bool check_flash_attention_hardware_support( c10::array_of( sycl::ext::oneapi::experimental::architecture::intel_gpu_pvc, sycl::ext::oneapi::experimental::architecture::intel_gpu_pvc_vg, - sycl::ext::oneapi::experimental::architecture::intel_gpu_bmg_g21); + sycl::ext::oneapi::experimental::architecture::intel_gpu_bmg_g21, + sycl::ext::oneapi::experimental::architecture::intel_gpu_bmg_g31); auto* device_prop = at::xpu::getCurrentDeviceProperties(); auto device_architecture = device_prop->architecture; @@ -39,7 +40,7 @@ bool check_flash_attention_hardware_support( device_architecture) == supported_architectures.end()) { if (debug) { TORCH_WARN( - "XPU device architecture does not support flash attention. Supported architectures are: intel_gpu_pvc, intel_gpu_pvc_vg, intel_gpu_bmg_g21."); + "XPU device architecture does not support flash attention. Supported architectures are: intel_gpu_pvc, intel_gpu_pvc_vg, intel_gpu_bmg_g21, intel_gpu_bmg_g31."); } return false; } @@ -80,9 +81,12 @@ inline bool check_flash_attention_datatype( inline bool check_flash_attention_head_dim_size( sdp_params const& params, bool debug) { - const int query_size_last = params.query.size(3); - const int key_size_last = params.key.size(3); - const int value_size_last = params.value.size(3); + // Use sym_size to preserve symbolic shapes during tracing. + // Using concrete .size() would materialize symbolic dimensions into static + // guards, preventing dynamic shape generalization across recompilations. + const auto query_size_last = params.query.sym_size(-1); + const auto key_size_last = params.key.sym_size(-1); + const auto value_size_last = params.value.sym_size(-1); const bool head_dims_equal = (query_size_last == key_size_last) && (query_size_last == value_size_last); @@ -101,7 +105,7 @@ inline bool check_flash_attention_head_dim_size( return false; } - constexpr auto max_supported_headdim = 192; + const auto max_supported_headdim = c10::SymInt(192); if (query_size_last > max_supported_headdim) { if (debug) { TORCH_WARN( diff --git a/aten/src/ATen/nnapi/CMakeLists.txt b/aten/src/ATen/nnapi/CMakeLists.txt index 9e367028e8a2d..4ba4048220370 100644 --- a/aten/src/ATen/nnapi/CMakeLists.txt +++ b/aten/src/ATen/nnapi/CMakeLists.txt @@ -3,7 +3,7 @@ if(PYTORCH_NNAPI_STANDALONE) cmake_minimum_required(VERSION 3.10 FATAL_ERROR) project(pytorch_nnapi) - set(CMAKE_CXX_STANDARD 14 CACHE STRING "The C++ standard whose features are requested to build this target.") + set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.") find_package(Torch REQUIRED) set(NNAPI_SRCS diff --git a/aten/src/ATen/test/native_test.cpp b/aten/src/ATen/test/native_test.cpp index 5f27ce4886e47..ff4507941a648 100644 --- a/aten/src/ATen/test/native_test.cpp +++ b/aten/src/ATen/test/native_test.cpp @@ -253,6 +253,9 @@ void test(TensorOptions T, TensorOptions AccT) { } TEST(TestNative, NativeTestCPU) { +#if defined(__aarch64__) + GTEST_SKIP() << "Known failure on AArch64 (label is too far / stack test mismatch)"; +#endif manual_seed(123); test(at::device(kCPU).dtype(kFloat), diff --git a/aten/src/ATen/test/test_install/CMakeLists.txt b/aten/src/ATen/test/test_install/CMakeLists.txt index c2d54feeb37d0..10d6805a092cb 100644 --- a/aten/src/ATen/test/test_install/CMakeLists.txt +++ b/aten/src/ATen/test/test_install/CMakeLists.txt @@ -2,9 +2,6 @@ cmake_minimum_required(VERSION 3.10) find_package(ATen REQUIRED) include_directories(${ATEN_INCLUDE_DIR}) -# C++17 -if(not MSVC) - set(CMAKE_CXX_FLAGS "--std=c++17 ${CMAKE_CXX_FLAGS}") -endif() +set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.") add_executable(main main.cpp) target_link_libraries(main ${ATEN_LIBRARIES}) diff --git a/aten/src/ATen/test/vitals.cpp b/aten/src/ATen/test/vitals.cpp deleted file mode 100644 index eaf1cc152bc37..0000000000000 --- a/aten/src/ATen/test/vitals.cpp +++ /dev/null @@ -1,96 +0,0 @@ -#include -#include - -#include -#include -#include -#include -#include - -using namespace at::vitals; -using ::testing::HasSubstr; - -TEST(Vitals, Basic) { - std::stringstream buffer; - - std::streambuf* sbuf = std::cout.rdbuf(); - std::cout.rdbuf(buffer.rdbuf()); - { - c10::utils::set_env("TORCH_VITAL", "1"); - TORCH_VITAL_DEFINE(Testing); - TORCH_VITAL(Testing, Attribute0) << 1; - TORCH_VITAL(Testing, Attribute1) << '1'; - TORCH_VITAL(Testing, Attribute2) << 1.0f; - TORCH_VITAL(Testing, Attribute3) << 1.0; - auto t = at::ones({1, 1}); - TORCH_VITAL(Testing, Attribute4) << t; - } - std::cout.rdbuf(sbuf); - - auto s = buffer.str(); - ASSERT_THAT(s, HasSubstr("Testing.Attribute0\t\t 1")); - ASSERT_THAT(s, HasSubstr("Testing.Attribute1\t\t 1")); - ASSERT_THAT(s, HasSubstr("Testing.Attribute2\t\t 1")); - ASSERT_THAT(s, HasSubstr("Testing.Attribute3\t\t 1")); - ASSERT_THAT(s, HasSubstr("Testing.Attribute4\t\t 1")); -} - -TEST(Vitals, MultiString) { - std::stringstream buffer; - - std::streambuf* sbuf = std::cout.rdbuf(); - std::cout.rdbuf(buffer.rdbuf()); - { - c10::utils::set_env("TORCH_VITAL", "1"); - TORCH_VITAL_DEFINE(Testing); - TORCH_VITAL(Testing, Attribute0) << 1 << " of " << 2; - TORCH_VITAL(Testing, Attribute1) << 1; - TORCH_VITAL(Testing, Attribute1) << " of "; - TORCH_VITAL(Testing, Attribute1) << 2; - } - std::cout.rdbuf(sbuf); - - auto s = buffer.str(); - ASSERT_THAT(s, HasSubstr("Testing.Attribute0\t\t 1 of 2")); - ASSERT_THAT(s, HasSubstr("Testing.Attribute1\t\t 1 of 2")); -} - -TEST(Vitals, OnAndOff) { - for (const auto i : c10::irange(2)) { - std::stringstream buffer; - - std::streambuf* sbuf = std::cout.rdbuf(); - std::cout.rdbuf(buffer.rdbuf()); - { - c10::utils::set_env("TORCH_VITAL", i ? "1" : "0"); - TORCH_VITAL_DEFINE(Testing); - TORCH_VITAL(Testing, Attribute0) << 1; - } - std::cout.rdbuf(sbuf); - - auto s = buffer.str(); - auto f = s.find("Testing.Attribute0\t\t 1"); - if (i) { - ASSERT_TRUE(f != std::string::npos); - } else { - ASSERT_TRUE(f == std::string::npos); - } - } -} - -TEST(Vitals, APIVitals) { - std::stringstream buffer; - bool rvalue = false; - std::streambuf* sbuf = std::cout.rdbuf(); - std::cout.rdbuf(buffer.rdbuf()); - { - c10::utils::set_env("TORCH_VITAL", "1"); - APIVitals api_vitals; - rvalue = api_vitals.setVital("TestingSetVital", "TestAttr", "TestValue"); - } - std::cout.rdbuf(sbuf); - - auto s = buffer.str(); - ASSERT_TRUE(rvalue); - ASSERT_THAT(s, HasSubstr("TestingSetVital.TestAttr\t\t TestValue")); -} diff --git a/aten/src/ATen/xpu/CachingHostAllocator.cpp b/aten/src/ATen/xpu/CachingHostAllocator.cpp index 2be5eaf81bb03..c6a1a807801e2 100644 --- a/aten/src/ATen/xpu/CachingHostAllocator.cpp +++ b/aten/src/ATen/xpu/CachingHostAllocator.cpp @@ -43,8 +43,7 @@ struct XPUCachingHostAllocatorImpl } bool stream_is_capturing(XPUStream s) const override { - return c10::xpu::CaptureStatus(s.queue().ext_oneapi_get_state()) == - c10::xpu::CaptureStatus::Recording; + return s.is_capturing(); } }; diff --git a/aten/src/ATen/xpu/XPUGeneratorImpl.cpp b/aten/src/ATen/xpu/XPUGeneratorImpl.cpp index 09c84eeef9cea..133a9e4728dd0 100644 --- a/aten/src/ATen/xpu/XPUGeneratorImpl.cpp +++ b/aten/src/ATen/xpu/XPUGeneratorImpl.cpp @@ -97,7 +97,7 @@ void XPUGeneratorState::increase(uint64_t increment) { } // State can be used by multiple graph -void XPUGeneratorState::register_graph(xpu::XPUGraph* graph) { +void XPUGeneratorState::register_graph(xpu::XPUGraphImpl* graph) { // Ensures that the RNG state is not currently being captured. at::xpu::assertNotCapturing( "Cannot register the state during capturing stage."); @@ -113,7 +113,7 @@ void XPUGeneratorState::register_graph(xpu::XPUGraph* graph) { } } -void XPUGeneratorState::unregister_graph(xpu::XPUGraph* graph) { +void XPUGeneratorState::unregister_graph(xpu::XPUGraphImpl* graph) { TORCH_CHECK( registered_graphs_.find(graph) != registered_graphs_.end(), "The graph should be registered to the state"); @@ -280,12 +280,12 @@ uint64_t XPUGeneratorImpl::philox_offset_per_thread() const { } } -void XPUGeneratorImpl::register_graph(xpu::XPUGraph* graph) { +void XPUGeneratorImpl::register_graph(xpu::XPUGraphImpl* graph) { graph->register_generator_state(state_); state_->register_graph(graph); } -void XPUGeneratorImpl::unregister_graph(xpu::XPUGraph* graph) { +void XPUGeneratorImpl::unregister_graph(xpu::XPUGraphImpl* graph) { state_->unregister_graph(graph); } diff --git a/aten/src/ATen/xpu/XPUGeneratorImpl.h b/aten/src/ATen/xpu/XPUGeneratorImpl.h index 8ee4967a4f13f..853d1f28f006a 100644 --- a/aten/src/ATen/xpu/XPUGeneratorImpl.h +++ b/aten/src/ATen/xpu/XPUGeneratorImpl.h @@ -8,7 +8,7 @@ namespace at { namespace xpu { -struct XPUGraph; +struct XPUGraphImpl; } struct XPUGeneratorState : public c10::intrusive_ptr_target { @@ -16,7 +16,7 @@ struct XPUGeneratorState : public c10::intrusive_ptr_target { uint64_t philox_offset_per_thread_; uint32_t offset_intragraph_; bool capturing_{}; - std::unordered_set registered_graphs_; + std::unordered_set registered_graphs_; at::TensorBase seed_extragraph_{}; at::TensorBase offset_extragraph_{}; @@ -29,8 +29,8 @@ struct XPUGeneratorState : public c10::intrusive_ptr_target { offset_intragraph_(offset_intragraph) {} void increase(uint64_t increment); - void register_graph(xpu::XPUGraph* graph); - void unregister_graph(xpu::XPUGraph* graph); + void register_graph(xpu::XPUGraphImpl* graph); + void unregister_graph(xpu::XPUGraphImpl* graph); void capture_prologue(); uint64_t capture_epilogue(); void replay_prologue(uint64_t wholegraph_increment); @@ -62,8 +62,8 @@ struct TORCH_XPU_API XPUGeneratorImpl : public GeneratorImpl { void set_philox_offset_per_thread(uint64_t offset); uint64_t philox_offset_per_thread() const; - void register_graph(xpu::XPUGraph* graph); - void unregister_graph(xpu::XPUGraph* graph); + void register_graph(xpu::XPUGraphImpl* graph); + void unregister_graph(xpu::XPUGraphImpl* graph); PhiloxXpuState philox_xpu_state(uint64_t increment); std::pair philox_engine_inputs(uint64_t increment); static c10::DeviceType device_type(); diff --git a/aten/src/ATen/xpu/XPUGraph.cpp b/aten/src/ATen/xpu/XPUGraph.cpp index 3d3a22e5b7219..d6d3cad494a04 100644 --- a/aten/src/ATen/xpu/XPUGraph.cpp +++ b/aten/src/ATen/xpu/XPUGraph.cpp @@ -1,6 +1,5 @@ #include #include -#include #include #include @@ -16,23 +15,43 @@ MempoolId_t graph_pool_handle() { return c10::xpu::MemPool::graph_pool_handle(); } -XPUGraph::XPUGraph(bool keep_graph) +XPUGraphImpl::XPUGraphImpl(const GraphImplArgs& args) : capture_stream_(at::xpu::getCurrentXPUStream()), - keep_graph_(keep_graph) {} + keep_graph_(args.keep_graph) {} -void XPUGraph::register_generator_state( +void XPUGraphImpl::register_generator_state( c10::intrusive_ptr state) { captured_generator_states_[std::move(state)] = 0; } -void XPUGraph::register_generator_state(const at::Generator& generator) { +void XPUGraphImpl::register_generator_state(const at::Generator& generator) { c10::intrusive_ptr xpu_gen = dynamic_intrusive_pointer_cast( generator.getIntrusivePtr()); xpu_gen->register_graph(this); } -void XPUGraph::capture_begin(MempoolId_t pool) { +void XPUGraphImpl::capture_begin( + MempoolId_t pool /*={0,0}*/, + GraphCaptureMode capture_mode) { + switch (capture_mode) { + case GraphCaptureMode::Default: + break; + + case GraphCaptureMode::Global: + case GraphCaptureMode::ThreadLocal: + case GraphCaptureMode::Relaxed: + TORCH_WARN( + "XPUGraph currently only support default GraphCaptureMode. " + "Falling back to default capture behavior."); + break; + + default: + TORCH_CHECK( + false, + "Invalid GraphCaptureMode value: ", + static_cast(capture_mode)); + } TORCH_CHECK( !has_graph_exec_, "This XPUGraph instance already owns a captured graph. " @@ -90,7 +109,7 @@ void XPUGraph::capture_begin(MempoolId_t pool) { capture_stream_.queue().ext_oneapi_get_state() == queue_state::recording); } -void XPUGraph::capture_end() { +void XPUGraphImpl::capture_end() { auto stream = at::xpu::getCurrentXPUStream(); TORCH_CHECK( @@ -125,7 +144,7 @@ void XPUGraph::capture_end() { } } -void XPUGraph::instantiate() { +void XPUGraphImpl::instantiate() { TORCH_CHECK( capture_ended_, "capture_end() must have been called before calling instantiate"); @@ -141,7 +160,7 @@ void XPUGraph::instantiate() { has_graph_exec_ = true; } -void XPUGraph::replay() { +void XPUGraphImpl::replay() { TORCH_CHECK( capture_ended_, "Called XPUGraph::replay without a preceding successful capture."); @@ -164,7 +183,7 @@ void XPUGraph::replay() { queue.ext_oneapi_graph(*graph_exec_); } -void XPUGraph::reset() { +void XPUGraphImpl::reset() { if (capture_ended_) { c10::xpu::XPUCachingAllocator::releasePool(capture_dev_, mempool_id_); at::getHostAllocator(at::kXPU)->release_pool(mempool_id_); @@ -180,11 +199,11 @@ void XPUGraph::reset() { } } -void XPUGraph::enable_debug_mode() { +void XPUGraphImpl::enable_debug_mode() { _xpu_graphs_debug = true; } -void XPUGraph::debug_dump(const std::string& debug_path) { +void XPUGraphImpl::debug_dump(const std::string& debug_path) { TORCH_CHECK( debug_path.size() >= 4 && debug_path.substr(debug_path.size() - 4) == ".dot", @@ -207,7 +226,7 @@ void XPUGraph::debug_dump(const std::string& debug_path) { } } -xpuGraph_t* XPUGraph::raw_xpu_graph() { +xpuGraph_t* XPUGraphImpl::raw_xpu_graph() { TORCH_CHECK( keep_graph_, "You cannot access the raw xpuGraph_t instance unless XPUGraph was initialized with keep_graph=true"); @@ -217,7 +236,7 @@ xpuGraph_t* XPUGraph::raw_xpu_graph() { return graph_.get(); } -xpuGraphExec_t* XPUGraph::raw_xpu_graph_exec() { +xpuGraphExec_t* XPUGraphImpl::raw_xpu_graph_exec() { TORCH_CHECK( has_graph_exec_, "You cannot access the raw xpuGraphExec_t instance until instantiate() has been called"); @@ -226,14 +245,14 @@ xpuGraphExec_t* XPUGraph::raw_xpu_graph_exec() { // Returns an id another graph's capture_begin can use to share the same memory // pool as this graph. -MempoolId_t XPUGraph::pool() { +MempoolId_t XPUGraphImpl::pool() const { TORCH_CHECK( capture_ended_, "Called XPUGraph::pool() without a preceding successful capture."); return mempool_id_; } -XPUGraph::~XPUGraph() { +XPUGraphImpl::~XPUGraphImpl() { for (auto& [generator_state, wholegraph_increments] : captured_generator_states_) { generator_state->unregister_graph(this); @@ -241,4 +260,6 @@ XPUGraph::~XPUGraph() { reset(); } +REGISTER_GRAPH_IMPL(XPU, XPUGraphImpl) + } // namespace at::xpu diff --git a/aten/src/ATen/xpu/XPUGraph.h b/aten/src/ATen/xpu/XPUGraph.h index 0f23a1e093ee0..fd60bb5b6ba9e 100644 --- a/aten/src/ATen/xpu/XPUGraph.h +++ b/aten/src/ATen/xpu/XPUGraph.h @@ -1,18 +1,15 @@ #pragma once #include +#include +#include #include #include #include #include #include -namespace at { - -struct Generator; -struct XPUGeneratorState; - -namespace xpu { +namespace at::xpu { TORCH_XPU_API MempoolId_t graph_pool_handle(); @@ -21,21 +18,26 @@ using xpuGraph_t = sycl::ext::oneapi::experimental::command_graph< using xpuGraphExec_t = sycl::ext::oneapi::experimental::command_graph< sycl::ext::oneapi::experimental::graph_state::executable>; -struct TORCH_XPU_API XPUGraph { - XPUGraph(bool keep_graph = false); - ~XPUGraph(); +struct TORCH_XPU_API XPUGraphImpl : public at::GraphImplInterface { + XPUGraphImpl(const GraphImplArgs& args = {}); + ~XPUGraphImpl() override; + + C10_DISABLE_COPY_AND_ASSIGN(XPUGraphImpl); void register_generator_state( c10::intrusive_ptr state); void register_generator_state(const at::Generator& generator); - void capture_begin(MempoolId_t pool = {0, 0}); - void capture_end(); - void instantiate(); - void replay(); - void reset(); - MempoolId_t pool(); - void enable_debug_mode(); - void debug_dump(const std::string& debug_path); + + void capture_begin( + MempoolId_t pool = {0, 0}, + GraphCaptureMode capture_mode = GraphCaptureMode::Default) override; + void capture_end() override; + void instantiate() override; + void replay() override; + void reset() override; + MempoolId_t pool() const override; + void enable_debug_mode() override; + void debug_dump(const std::string& debug_path) override; xpuGraph_t* raw_xpu_graph(); xpuGraphExec_t* raw_xpu_graph_exec(); @@ -59,5 +61,58 @@ struct TORCH_XPU_API XPUGraph { bool keep_graph_; }; -} // namespace xpu -} // namespace at +struct TORCH_XPU_API XPUGraph { + XPUGraph(bool keep_graph = false) { + GraphImplArgs args; + args.keep_graph = keep_graph; + impl_ = std::make_unique(args); + } + ~XPUGraph() = default; + + C10_DISABLE_COPY_AND_ASSIGN(XPUGraph); + XPUGraph(XPUGraph&& other) = delete; + XPUGraph& operator=(XPUGraph&& other) = delete; + + void register_generator_state( + c10::intrusive_ptr state) { + impl_->register_generator_state(state); + } + void register_generator_state(const at::Generator& generator) { + impl_->register_generator_state(generator); + } + void capture_begin(MempoolId_t pool = {0, 0}) { + impl_->capture_begin(pool); + } + void capture_end() { + impl_->capture_end(); + } + void instantiate() { + impl_->instantiate(); + } + void replay() { + impl_->replay(); + } + void reset() { + impl_->reset(); + } + MempoolId_t pool() const { + return impl_->pool(); + } + void enable_debug_mode() { + impl_->enable_debug_mode(); + } + void debug_dump(const std::string& debug_path) { + impl_->debug_dump(debug_path); + } + xpuGraph_t* raw_xpu_graph() { + return impl_->raw_xpu_graph(); + } + xpuGraphExec_t* raw_xpu_graph_exec() { + return impl_->raw_xpu_graph_exec(); + } + + private: + std::unique_ptr impl_; +}; + +} // namespace at::xpu diff --git a/aten/tools/run_tests.sh b/aten/tools/run_tests.sh index 93e48aec90851..2c57d428f58ea 100755 --- a/aten/tools/run_tests.sh +++ b/aten/tools/run_tests.sh @@ -55,6 +55,7 @@ run_if_exists cuda_allocator_test if [ "$VALGRIND" == "ON" ]; then # NB: As these tests are invoked by valgrind, let's leave them for now as it's # unclear if valgrind -> python -> gtest would work + export LD_LIBRARY_PATH="${CPP_TESTS_DIR}:${LD_LIBRARY_PATH}" valgrind --suppressions="$VALGRIND_SUP" --error-exitcode=1 "${CPP_TESTS_DIR}/basic" --gtest_filter='-*CUDA' if [[ -x ${CPP_TESTS_DIR}/tensor_interop_test ]]; then valgrind --suppressions="$VALGRIND_SUP" --error-exitcode=1 "${CPP_TESTS_DIR}/tensor_interop_test" diff --git a/aten/tools/valgrind.sup b/aten/tools/valgrind.sup index 585487c4d2be2..c86c374bd1307 100644 --- a/aten/tools/valgrind.sup +++ b/aten/tools/valgrind.sup @@ -10,6 +10,16 @@ ... } +{ + + Memcheck:Addr8 + fun:strncmp + fun:is_dst + ... + fun:decompose_rpath + ... +} + { ignore_empty_generic_uninitialised_conditional_jump Memcheck:Cond diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_torchbench_inference.csv index 0799b804bbf8c..8fd420c017d12 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_torchbench_inference.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_amp_freezing_torchbench_inference.csv @@ -50,47 +50,47 @@ densenet121,pass,0 -detectron2_fasterrcnn_r_101_c4,fail_accuracy,42 +detectron2_fasterrcnn_r_101_c4,fail_accuracy,43 -detectron2_fasterrcnn_r_101_dc5,fail_accuracy,42 +detectron2_fasterrcnn_r_101_dc5,fail_accuracy,43 -detectron2_fasterrcnn_r_101_fpn,fail_accuracy,46 +detectron2_fasterrcnn_r_101_fpn,fail_accuracy,47 -detectron2_fasterrcnn_r_50_c4,fail_accuracy,42 +detectron2_fasterrcnn_r_50_c4,fail_accuracy,43 -detectron2_fasterrcnn_r_50_dc5,fail_accuracy,42 +detectron2_fasterrcnn_r_50_dc5,fail_accuracy,43 -detectron2_fasterrcnn_r_50_fpn,fail_accuracy,46 +detectron2_fasterrcnn_r_50_fpn,fail_accuracy,47 -detectron2_fcos_r_50_fpn,fail_accuracy,24 +detectron2_fcos_r_50_fpn,fail_accuracy,25 -detectron2_maskrcnn_r_101_c4,fail_accuracy,56 +detectron2_maskrcnn_r_101_c4,fail_accuracy,57 -detectron2_maskrcnn_r_101_fpn,fail_accuracy,62 +detectron2_maskrcnn_r_101_fpn,fail_accuracy,63 -detectron2_maskrcnn_r_50_c4,fail_accuracy,56 +detectron2_maskrcnn_r_50_c4,fail_accuracy,57 -detectron2_maskrcnn_r_50_fpn,fail_accuracy,62 +detectron2_maskrcnn_r_50_fpn,fail_accuracy,63 @@ -242,7 +242,7 @@ vgg16,pass,0 -vision_maskrcnn,fail_accuracy,30 +vision_maskrcnn,fail_accuracy,31 diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_huggingface_inference.csv index e4aabce10466d..2712e2d1a5012 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_huggingface_inference.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_huggingface_inference.csv @@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,0 -DistillGPT2,pass,2 +DistillGPT2,pass,3 @@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0 -M2M100ForConditionalGeneration,pass,7 +M2M100ForConditionalGeneration,pass,0 diff --git a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv index 8ce784146dc54..0a51409c04ac3 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/cpu_inductor_torchbench_inference.csv @@ -82,15 +82,15 @@ detectron2_maskrcnn_r_101_c4,fail_accuracy,56 -detectron2_maskrcnn_r_101_fpn,pass,63 +detectron2_maskrcnn_r_101_fpn,pass,62 -detectron2_maskrcnn_r_50_c4,pass,57 +detectron2_maskrcnn_r_50_c4,pass,56 -detectron2_maskrcnn_r_50_fpn,pass,63 +detectron2_maskrcnn_r_50_fpn,pass,62 diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_huggingface_inference.csv index e4aabce10466d..2712e2d1a5012 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_huggingface_inference.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_cpu_inductor_huggingface_inference.csv @@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,0 -DistillGPT2,pass,2 +DistillGPT2,pass,3 @@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0 -M2M100ForConditionalGeneration,pass,7 +M2M100ForConditionalGeneration,pass,0 diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv index 72fd3af5beeda..3ef5716e9306a 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_inference.csv @@ -74,7 +74,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0 -detectron2_fcos_r_50_fpn,pass,22 +detectron2_fcos_r_50_fpn,pass,23 diff --git a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_training.csv index e4c20cfebf465..9dd6d40930389 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_training.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/dynamic_inductor_torchbench_training.csv @@ -190,7 +190,7 @@ vgg16,pass,6 -vision_maskrcnn,fail_to_run,37 +vision_maskrcnn,fail_to_run,38 diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_inference.csv index 87dd88078f222..1333ff866469e 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_inference.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_inference.csv @@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,0 -DistillGPT2,pass,2 +DistillGPT2,pass,3 @@ -50,7 +50,7 @@ LayoutLMForMaskedLM,pass,0 -M2M100ForConditionalGeneration,pass,7 +M2M100ForConditionalGeneration,pass,0 @@ -122,4 +122,4 @@ meta-llama/Llama-3.2-1B,pass,0 -openai/whisper-tiny,pass,5 +openai/whisper-tiny,pass,7 diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_training.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_training.csv index 5ca03b5ecf9fb..452bf8707fbfb 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_training.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_training.csv @@ -10,7 +10,7 @@ AllenaiLongformerBase,pass,9 -BartForCausalLM,pass,6 +BartForCausalLM,pass,7 @@ -30,7 +30,7 @@ DistilBertForMaskedLM,pass,5 -DistillGPT2,pass,7 +DistillGPT2,pass,8 @@ -50,11 +50,11 @@ LayoutLMForMaskedLM,pass,5 -M2M100ForConditionalGeneration,pass,11 +M2M100ForConditionalGeneration,pass,4 -MBartForCausalLM,pass,6 +MBartForCausalLM,pass,7 @@ -70,15 +70,15 @@ MobileBertForMaskedLM,pass,3 -OPTForCausalLM,pass,8 +OPTForCausalLM,pass,7 -PLBartForCausalLM,pass,6 +PLBartForCausalLM,pass,7 -PegasusForCausalLM,pass,6 +PegasusForCausalLM,pass,7 diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv index bc98e325ec784..d04a273049dbf 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv @@ -74,7 +74,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0 -detectron2_fcos_r_50_fpn,pass,22 +detectron2_fcos_r_50_fpn,pass,23 diff --git a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv index e4c20cfebf465..90c1f0972a3b2 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv @@ -170,6 +170,10 @@ soft_actor_critic,pass,6 +speech_transformer,pass,16 + + + squeezenet1_1,pass,6 @@ -190,7 +194,7 @@ vgg16,pass,6 -vision_maskrcnn,fail_to_run,37 +vision_maskrcnn,fail_to_run,38 diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_inference.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_inference.csv index d1150f849e2ee..0c22866fffccb 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_inference.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_inference.csv @@ -74,7 +74,7 @@ detectron2_fasterrcnn_r_50_fpn,eager_fail_to_run,0 -detectron2_fcos_r_50_fpn,pass,22 +detectron2_fcos_r_50_fpn,pass,23 diff --git a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_training.csv b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_training.csv index 877277a5aa192..a4d42c7d7723a 100644 --- a/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_training.csv +++ b/benchmarks/dynamo/ci_expected_accuracy/rocm/dynamic_inductor_torchbench_training.csv @@ -90,7 +90,7 @@ microbench_unbacked_tolist_sum,pass,9 -mnasnet1_0,pass,0 +mnasnet1_0,pass,7 @@ -102,7 +102,7 @@ mobilenet_v2_quantized_qat,eager_fail_to_run,0 -mobilenet_v3_large,pass,0 +mobilenet_v3_large,pass,7 @@ -141,7 +141,7 @@ pytorch_unet,pass_due_to_skip,7 -resnet152,eager_two_runs_differ,0 +resnet152,eager_two_runs_differ,7 @@ -193,7 +193,7 @@ vgg16,pass,6 -vision_maskrcnn,fail_to_run,37 +vision_maskrcnn,fail_to_run,38 diff --git a/benchmarks/dynamo/common.py b/benchmarks/dynamo/common.py index 913291ebe0d3e..eb3b44dfdc918 100644 --- a/benchmarks/dynamo/common.py +++ b/benchmarks/dynamo/common.py @@ -2024,17 +2024,18 @@ def cast_based_on_args(self, model, example_inputs): return model, example_inputs - def validate_model(self, model, example_inputs): + def validate_model(self, name, model, example_inputs): """ Runs the eager model with example inputs to ensure that eager passes. """ model = self.deepcopy_model(model) example_inputs = clone_inputs(example_inputs) model, example_inputs = self.cast_based_on_args(model, example_inputs) - try: - self.model_iter_fn(model, example_inputs) - except Exception as e: - raise RuntimeError("Eager run failed") from e + with self.pick_grad(name, self.args.training): + try: + self.model_iter_fn(model, example_inputs) + except Exception as e: + raise RuntimeError("Eager run failed") from e def maybe_cast(self, model, example_inputs): model, example_inputs = self.cast_based_on_args(model, example_inputs) @@ -3204,6 +3205,11 @@ def parse_args(args=None): parser.add_argument( "--freezing", action="store_true", help="turn on freezing", default=False ) + parser.add_argument( + "--deterministic", + action="store_true", + help="Enable deterministic mode (torch.use_deterministic_algorithms, cudnn.deterministic, etc.)", + ) parser.add_argument( "--inductor-config", "-c", @@ -3867,7 +3873,8 @@ def print_comparison(all_results): ) print(f"{'=' * 80}", flush=True) - # Build base command, stripping --compare-backed-unbacked and --only + value + # Build base command, stripping --compare-backed-unbacked, --only, --filter and their values + # Handles both space-separated (--filter VALUE) and equals-separated (--filter=VALUE) forms filtered = [] skip_next = False for a in sys.argv: @@ -3876,8 +3883,13 @@ def print_comparison(all_results): if skip_next: skip_next = False continue - if a == "--only": - skip_next = True + if a == "--only" or a.startswith("--only="): + if "=" not in a: + skip_next = True + continue + if a == "--filter" or a.startswith("--filter="): + if "=" not in a: + skip_next = True continue filtered.append(a) base_cmd = [sys.executable, "-B"] + filtered @@ -3984,7 +3996,7 @@ def write_csv_when_exception(args, name: str, status: str, device=None): write_outputs(output_filename, headers, row) -def setup_determinism_for_accuracy_test(args): +def setup_determinism(args): if args.only is not None and args.only not in { "alexnet", "Background_Matting", @@ -4020,6 +4032,11 @@ def setup_determinism_for_accuracy_test(args): torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False torch.backends.mkldnn.deterministic = True + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + torch.backends.cudnn.allow_tf32 = False + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(False) + patch_torch_manual_seed() def run(runner, args, original_dir=None): @@ -4080,6 +4097,9 @@ def run(runner, args, original_dir=None): "DLRM+DDP is unsupported as it requires sharding the embedding layer separately from DDP" ) return sys.exit(-1) + if args.deterministic and not args.accuracy: + setup_determinism(args) + if args.accuracy: # Use small batch size. We use >1 batch size to ensure we test # batch_norm type of operators that work on batch dims. @@ -4103,9 +4123,8 @@ def run(runner, args, original_dir=None): args.use_eval_mode = True inductor_config.fallback_random = True - setup_determinism_for_accuracy_test(args) + setup_determinism(args) - os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" if args.only is not None and args.only in { "nvidia_deeprecommender", }: @@ -4113,13 +4132,6 @@ def run(runner, args, original_dir=None): torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False - torch.backends.cudnn.allow_tf32 = False - torch.backends.cuda.matmul.allow_tf32 = False - torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(False) - - # Remove randomness when torch manual seed is called - patch_torch_manual_seed() - # Some models e.g. yolov3 assert batch size on n_gpus if "CUDA_VISIBLE_DEVICES" not in os.environ and not args.multiprocess: args.device_index = "0" @@ -4331,8 +4343,8 @@ def run(runner, args, original_dir=None): ) model_iter_fn = baseline_ctx(runner.model_iter_fn) - # needed to avoid error that causes inconsistent timing due to: - # Unable to hit fast path of CUDAGraphs because of pending, uninvoked backwards + # needed to avoid CUDAGraph fast-path warning / inconsistent timing when prior + # outputs still require backward (see torch._inductor.cudagraph_trees) def model_iter_fn_and_mark_step(*args, **kwargs): torch.compiler.cudagraph_mark_step_begin() model_iter_fn(*args, **kwargs) diff --git a/benchmarks/dynamo/huggingface.py b/benchmarks/dynamo/huggingface.py index b91777f70ed98..d62d99bfc4ffb 100755 --- a/benchmarks/dynamo/huggingface.py +++ b/benchmarks/dynamo/huggingface.py @@ -489,7 +489,7 @@ def generate(self, _, example_inputs, collect_outputs=True): else: model.eval() - self.validate_model(model, example_inputs) + self.validate_model(model_name, model, example_inputs) return device, model_name, model, example_inputs, batch_size def iter_model_names(self, args): @@ -552,7 +552,7 @@ def compute_loss(self, pred): return pred[0] def forward_pass(self, mod, inputs, collect_outputs=True): - with self.autocast(**self.autocast_arg): + with torch.no_grad(), self.autocast(**self.autocast_arg): res = mod(**inputs) return res.logits if self.hf_llm else res diff --git a/benchmarks/dynamo/launch_compile_op_numerics.py b/benchmarks/dynamo/launch_compile_op_numerics.py new file mode 100644 index 0000000000000..0b6871b1c08bf --- /dev/null +++ b/benchmarks/dynamo/launch_compile_op_numerics.py @@ -0,0 +1,974 @@ +import argparse +import asyncio +import datetime +import functools +import hashlib +import itertools +import json +import logging +import os +import re +import shlex +import sys +from dataclasses import dataclass, field +from pathlib import Path + + +log = logging.getLogger(__name__) + +try: + import torch +except ImportError: + torch = None + +PYTORCH_NIGHTLY_CUDA_VERSIONS = ["12.6", "12.8", "13.0"] +PYTORCH_CUDA_VERSIONS = { + "2.9.1": ["12.6", "12.8", "13.0"], + "2.9.0": ["12.6", "12.8", "13.0"], + "2.8.0": ["12.6", "12.8", "12.9"], + "2.7.1": ["11.8", "12.6", "12.8"], + "2.7.0": ["11.8", "12.6", "12.8"], + "2.6.0": ["11.8", "12.4", "12.6"], + "2.5.1": ["11.8", "12.1", "12.4"], + "2.5.0": ["11.8", "12.1", "12.4"], + "2.4.1": ["11.8", "12.1", "12.4"], + "2.4.0": ["11.8", "12.1", "12.4"], + "2.3.1": ["11.8", "12.1"], + "2.3.0": ["11.8", "12.1"], + "2.2.2": ["11.8", "12.1"], + "2.2.1": ["11.8", "12.1"], + "2.2.0": ["11.8", "12.1"], + "2.1.2": ["11.8", "12.1"], + "2.1.1": ["11.8", "12.1"], + "2.1.0": ["11.8", "12.1"], + "2.0.1": ["11.8"], + "2.0.0": ["11.8"], +} + +ENABLED_CONFIGS = [ + # ("git:5f09e6a6c93e0b5bf75b635cddc03b85bbe85938", "12.8"), + ("nightly", "13.0"), + # ("2.9.1", "12.8"), + # ("2.7.1", "12.8"), +] + +PYTHON_VERSION = "3.11" + + +@dataclass +class Mode: + compile: bool + backend: str | None = None + mode: str | None = None + options: dict | None = None + env: dict = field(default_factory=dict) + + def __call__(self, fn): + if self.compile: + kwargs = {} + if self.backend is not None: + kwargs["backend"] = self.backend + if self.mode is not None: + kwargs["mode"] = self.mode + if self.options is not None: + options = { + k: v + for k, v in self.options.items() + if k in torch._inductor.list_options() + } + kwargs["options"] = options + return torch.compile(fn, **kwargs) + assert self.backend is None # noqa: S101 + assert self.mode is None # noqa: S101 + assert self.options is None # noqa: S101 + return fn + + +MODES = { + "eager": Mode(compile=False), + "decomp": Mode(compile=True, backend="aot_eager_decomp_partition"), + "compile_numerics": Mode( + compile=True, + options={ + "emulate_precision_casts": True, + "use_fast_math": False, + "emulate_division_rounding": True, + "eager_numerics.division_rounding": True, + "eager_numerics.disable_ftz": True, + }, + ), + "compile": Mode(compile=True), +} + +DTYPES = ["float32", "float16", "bfloat16"] + + +async def run(command, capture_output=True, capture_stderr=True, stderr=None, **kwargs): + log.info("Running command: %s", command) + proc = await asyncio.create_subprocess_exec( + *shlex.split(command), + stdout=asyncio.subprocess.PIPE if capture_output else None, + stderr=asyncio.subprocess.STDOUT if capture_stderr else stderr, + stdin=asyncio.subprocess.DEVNULL, + **kwargs, + ) + if capture_output: + result = await proc.communicate() + result = result[0].decode("utf-8") + log.info("Command result: %s", result) + log.info("Return code: %s", proc.returncode) + return result + else: + await proc.wait() + log.info("Return code: %s", proc.returncode) + return + + +async def copy_file_to_remote(hostname, source, destination): + await run(f"scp {source} {hostname}:{destination}") + + +async def execute_on_remote(hostname, command): + await run(f"ssh {hostname} {command}", capture_output=False) + + +async def copy_results_from_remote(hostname, run_id): + await run(f"mkdir -p results_{run_id}/logs") + await run(f"scp {hostname}:/workspace/result_*.jsonl results_{run_id}/") + await run( + f"scp {hostname}:/workspace/logs.tar.gz results_{run_id}/logs_{hostname}.tar.gz" + ) + await run(f"mkdir -p results_{run_id}/logs") + await run( + f"tar xzf results_{run_id}/logs_{hostname}.tar.gz -C results_{run_id}/logs --strip 2" + ) + await run( + "bash -c 'for f in *.trace; do cd $f; uvx tlparse --no-browser *.log & cd ..; done; wait'", + cwd=f"results_{run_id}/logs", + ) + await run(f"rm results_{run_id}/logs_{hostname}.tar.gz") + await run(f"bash -c 'rm results_{run_id}/logs/*.trace/*.log'") + await run(f"bash -c 'rm results_{run_id}/logs/*.trace/tl_out/raw*'") + + +async def do_numerics_test(args, hostname, gpu, run_id): + await copy_file_to_remote(hostname, __file__, "/workspace/run.py") + await execute_on_remote(hostname, f"python3 /workspace/run.py --worker --gpu {gpu}") + await copy_results_from_remote(hostname, run_id) + + +async def reserve_gpu(gpu): + stdout = await run( + f"gpu-dev reserve -g 1 -h 24 -t {gpu} --ignore-no-persist --disk none" + ) + # the output has a string like this: SSH Command: ssh gpu-dev-1db65ec7 + match = re.search(r"SSH Command: ssh (.*)", stdout) + hostname = match.group(1) + log.info( + "Hostname: %s, Reservation ID: %s", + hostname, + hostname.removeprefix("gpu-dev-"), + ) + return hostname, hostname.removeprefix("gpu-dev-") + + +async def cancel_reservation(reservation_id): + await run(f"gpu-dev cancel {reservation_id}") + + +async def run_on_gpu(gpu, command): + log.info("Reserving GPU %s", gpu) + hostname, reservation_id = await reserve_gpu(gpu) + try: + log.info("Running command on %s", hostname) + await command(hostname=hostname) + except Exception: + log.exception("Error running command on %s", gpu) + finally: + log.info("Cancelling reservation %s", reservation_id) + await cancel_reservation(reservation_id) + + +async def launcher(args): + run_id = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + await asyncio.gather( + *[ + run_on_gpu( + gpu, + functools.partial(do_numerics_test, args=args, gpu=gpu, run_id=run_id), + ) + for gpu in args.gpu + ] + ) + await run(f"{sys.executable} {__file__} --html {run_id}") + + +def format_results_to_html(run_id): + results_dir = Path(f"results_{run_id}") + for f in results_dir.glob("*.jsonl"): + html_file = f.with_suffix(".html") + lines = f.read_text().splitlines(keepends=True) + data = [json.loads(line) for line in lines] + KEYS = data[0].keys() + print(KEYS) + FILTER_KEYS = [ + "gpu", + "pytorch_version", + "cuda_version", + "pytorch_mode", + "is_golden", + "data_type", + "function", + "pass_type", + "category", + "match_full", + "match_normal", + ] + KEY_VALUES = {k: sorted({line[k] for line in data}) for k in FILTER_KEYS} + with html_file.open("w") as html_f: + html_f.write("\n") + html_f.write(f"

Results for {f.stem}

\n") + for key, values in KEY_VALUES.items(): + html_f.write( + f"
{key} ({len(values)})\n" + ) + html_f.write( + f"Enable all\n" # noqa: B950 + ) + html_f.write( + f"Disable all\n" # noqa: B950 + ) + for value in values: + safe_value = str(value).replace(".", "_") + html_f.write( + f"\n" + ) + html_f.write( + f"\n" # noqa: B950 + ) + html_f.write("
") + # one more filter: show "first" mode only + html_f.write("
First mode that fails only\n") + html_f.write( + "\n" + ) + html_f.write("
\n") + # column filter: + html_f.write( + "
Column filter\n" + ) + html_f.write( + "Enable all\n" # noqa: B950 + ) + html_f.write( + "Disable all\n" # noqa: B950 + ) + DEFAULT_HIDDEN_KEYS = [ + "gpu", + "pytorch_version", + "cuda_version", + "is_golden", + ] + for key in KEYS: + html_f.write( + f"\n" # noqa: B950 + ) + html_f.write( + f"\n" + ) + html_f.write("
\n") + html_f.write("\n") + html_f.write("\n") + for key in KEYS: + html_f.write(f"\n") + + html_f.write("\n") + html_f.write("\n") + html_f.write("\n") + for line in lines: + data = json.loads(line) + classes = [ + f"visible-{k}-{str(data[k]).replace('.', '_')}" for k in KEY_VALUES + ] + html_f.write(f"\n") + for key in KEYS: + html_f.write(f"\n") + html_f.write( + f"" + ) + html_f.write( + f"\n" + ) + html_f.write("\n") + if data["mismatch_sample"]: + html_f.write( + f"\n" + ) + html_f.write(f"\n") + html_f.write("\n") + + html_f.write("
{key}logtlparse
\n") + if key == "mismatch_sample": + html_f.write( + f"\n" + ) + html_f.write( + f"\n" # noqa: B950 + ) + else: + html_f.write(f"{data[key]}\n") + html_f.write("logtrace
\n")
+                    html_f.write("\n")
+                    html_f.write("\n")
+                    html_f.write("\n")
+                    html_f.write("\n")
+                    html_f.write("\n")
+                    html_f.write("\n")
+                    html_f.write("\n")
+                    html_f.write("\n")
+                    for sample in data["mismatch_sample"]:
+                        html_f.write("\n")
+                        html_f.write(f"\n")
+                        html_f.write(
+                            f"\n"
+                        )
+                        html_f.write(f"\n")
+                        html_f.write(f"\n")
+                        html_f.write(f"\n")
+                        html_f.write("\n")
+                    html_f.write("
posinputoutputgoldenrel_err
{sample['pos']}{', '.join(map(str, sample['input']))}{sample['output']}{sample['golden']}{sample['rel_err']}
\n")
+                    html_f.write("
\n") + html_f.write("\n") + html_f.write("\n") + + +global_utils_lock = asyncio.Lock() +cuda_download_lock = asyncio.Lock() +venv_creation_lock = asyncio.Lock() +pytorch_build_lock = asyncio.Lock() +numerics_test_lock = asyncio.Lock() + + +async def create_results(args, config): + pytorch_version, cuda_version = config + async with global_utils_lock: + await ensure_global_utils() + async with cuda_download_lock: + await ensure_cuda(args, cuda_version) + async with venv_creation_lock: + venv = await create_venv(args, config) + async with pytorch_build_lock: + await maybe_build_pytorch(args, config, venv) + async with numerics_test_lock: + await run_numerics_test(args, config, venv) + + +async def ensure_global_utils(): + if os.environ.get("HAS_GLOBAL_UTILS", "0") == "1": + return + await run("mkdir global_bin") + os.environ["PATH"] = f"/workspace/global_bin:{os.environ.get('PATH', '')}" + + os.environ["UV_INSTALL_DIR"] = "/workspace/global_bin" + os.environ["UV_CACHE_DIR"] = "/workspace/.cache/uv" + await run("wget --no-verbose https://astral.sh/uv/install.sh") + await run("bash install.sh") + await run("rm install.sh") + + await run( + "wget --no-verbose https://github.com/seeraven/gitcache/releases/download/v1.0.29/gitcache_v1.0.29_Ubuntu22.04_x86_64" # noqa: B950 + ) + await run("mv gitcache_v1.0.29_Ubuntu22.04_x86_64 global_bin/gitcache") + await run("chmod +x global_bin/gitcache") + which_gitcache = await run("which gitcache") + await run(f"ln -s {which_gitcache.strip()} global_bin/git") + os.environ["GITCACHE_DIR"] = "/workspace/.cache/gitcache" + + await run("sudo apt-get update") + await run("sudo apt-get install -y ccache git-lfs") + os.environ["CMAKE_C_COMPILER_LAUNCHER"] = "ccache" + os.environ["CMAKE_CXX_COMPILER_LAUNCHER"] = "ccache" + os.environ["CMAKE_CUDA_COMPILER_LAUNCHER"] = "ccache" + os.environ["CCACHE_DIR"] = "/workspace/.cache/ccache" + os.environ["CCACHE_NOHASHDIR"] = "1" + os.environ["CCACHE_BASEDIR"] = "/workspace" + os.environ["CMAKE_BUILD_PARALLEL_LEVEL"] = "20" + os.environ["MAX_JOBS"] = "8" + os.environ["MAKEFLAGS"] = "-j20" + os.environ["NINJAFLAGS"] = "-j20" + + os.environ["HAS_GLOBAL_UTILS"] = "1" + + +async def ensure_cuda(args, cuda_version): + if os.path.exists(f"/workspace/cuda-{cuda_version}"): + return + await run("sudo apt-get update") + await run( + f"sudo apt-get install -y cuda-toolkit-{cuda_version.split('.')[0]}-{cuda_version.split('.')[1]}" + ) + + +async def create_venv(args, config): + pytorch_version, cuda_version = config + env = os.environ.copy() + # cuda_home = f"/usr/local/cuda-{cuda_version}" + cuda_home = ( + f"/usr/local/cuda-{cuda_version.split('.')[0]}.{cuda_version.split('.')[1]}" + ) + env["CUDA_HOME"] = cuda_home + env["CUDA_ROOT"] = cuda_home + env["PATH"] = f"{cuda_home}/bin:{env.get('PATH', '')}" + env["LD_LIBRARY_PATH"] = f"{cuda_home}/lib64:{env.get('LD_LIBRARY_PATH', '')}" + env["CUDA_NVCC_EXECUTABLE"] = f"{cuda_home}/bin/nvcc" + env["PYTORCH_NVCC"] = f"{cuda_home}/bin/nvcc" + venv_dir = f"/workspace/venv-pytorch_{pytorch_version}-cuda_{cuda_version}" + await run(f"mkdir -p {venv_dir}") + await run(f"uv venv -p {PYTHON_VERSION} --managed-python", cwd=venv_dir, env=env) + # update env with venv paths + env["VIRTUAL_ENV"] = f"{venv_dir}/.venv" + env["PATH"] = f"{venv_dir}/.venv/bin:{env.get('PATH', '')}" + await run("uv pip install pip numpy", env=env, cwd=venv_dir) + + index_url = f"https://download.pytorch.org/whl/cu{cuda_version.replace('.', '')}" + if pytorch_version == "nightly": + await run( + f"uv pip install --pre torch --index-url {index_url}", env=env, cwd=venv_dir + ) + elif not pytorch_version.startswith("git:"): + await run( + f"uv pip install torch=={pytorch_version} --index-url {index_url}", + env=env, + cwd=venv_dir, + ) + return {"env": env, "cwd": venv_dir} + + +async def maybe_build_pytorch(args, config, venv): + pytorch_version, cuda_version = config + if pytorch_version.startswith("git:"): + raise NotImplementedError("Building PyTorch from git is not supported yet") + return + + +async def run_numerics_test(args, config, venv): + pytorch_version, cuda_version = config + golden_flag = "" + if args.create_golden: + golden_flag = " --create-golden" + for mode in MODES: + cmd = ( + f"python {__file__} --runner --gpu {args.gpu[0]}" + f" --pytorch-version {pytorch_version}" + f" --cuda-version {cuda_version}" + f" --mode {mode}" + f" --golden /workspace/golden{golden_flag}" + ) + await run(cmd, capture_output=False, **venv) + golden_flag = "" + + +def sortable_config_key(config): + pytorch_version, cuda_version = config + cuda_version = tuple(map(int, cuda_version.split("."))) + + if pytorch_version == "nightly": + pytorch_version = (-1,) + elif pytorch_version.startswith("git:"): + pytorch_version = (-2,) + else: + pytorch_version = tuple(map(int, pytorch_version.split("."))) + + return (pytorch_version, cuda_version) + + +async def worker(args): + assert len(args.gpu) == 1 # noqa: S101 + # find golden task + # a config is a tuple (pytorch_version, cuda_version) + # where pytorch_version can either be a version string or nightly or a git hash + golden_config = sorted(ENABLED_CONFIGS, key=sortable_config_key)[-1] + await run("mkdir -p /workspace/logs") + args.create_golden = True + await create_results(args, golden_config) + args.create_golden = False + await asyncio.gather( + *[ + create_results(args, config) + for config in ENABLED_CONFIGS + if config != golden_config + ] + ) + await run("tar czf /workspace/logs.tar.gz /workspace/logs") + + +CATEGORIES = { + "reduction": ["torch.sum", "torch.mean", "torch.softmax"], + "matrix": ["torch.matmul"], + "normalization": [ + "torch.nn.functional.layer_norm", + "torch.nn.functional.rms_norm", + ], + "activation": [ + "torch.nn.functional.relu", + "torch.nn.functional.sigmoid", + "torch.nn.functional.tanh", + "torch.nn.functional.gelu", + "torch.nn.functional.silu", + ], + "elementary": [ + "torch.sin", + "torch.cos", + "torch.tan", + "torch.sigmoid", + "torch.exp", + "torch.exp2", + "torch.log", + "torch.log2", + "torch.sqrt", + "torch.erf", + "torch.reciprocal", + "torch.rsqrt", + ], + "binary": [ + "torch.add", + "torch.sub", + "torch.mul", + "torch.div", + "torch.pow", + ], +} + +PASS_TYPES = { + "reduction": ["fwd", "bwd_0"], + "matrix": ["fwd", "bwd_0", "bwd_1"], + "normalization": ["fwd", "bwd_0", "bwd_1"], + "activation": ["fwd", "bwd_0"], + "elementary": ["fwd", "bwd_0"], + "binary": ["fwd", "bwd_0", "bwd_1"], +} + +PASSES = {} + + +def register_pass(pass_type): + def decorator(func): + PASSES[pass_type] = func + return func + + return decorator + + +@register_pass("fwd") +def pass_fwd(callable): + def wrapper(*args): + return callable(*args) + + return wrapper + + +@register_pass("bwd_0") +def pass_bwd_0(callable): + def wrapper(*args): + args = [arg.detach().requires_grad_(idx == 0) for idx, arg in enumerate(args)] + output = callable(*args) + output.sum().backward() + return args[0].grad.detach() + + return wrapper + + +@register_pass("bwd_1") +def pass_bwd_1(callable): + def wrapper(*args): + args = [arg.detach().requires_grad_(idx == 1) for idx, arg in enumerate(args)] + output = callable(*args) + output.sum().backward() + return args[1].grad.detach() + + return wrapper + + +def filter_nan(a): + return torch.where(torch.isnan(a), 0, a) + + +def mantissa_bits(dtype) -> int: + return {torch.float32: 23, torch.bfloat16: 5, torch.float16: 10}[dtype] + + +def float_to_int_type(dtype): + return { + torch.float32: torch.uint32, + torch.bfloat16: torch.uint16, + torch.float16: torch.uint16, + }[dtype] + + +def generate_test_tensor(dtype, slow=False): + if dtype in [torch.float16, torch.bfloat16]: + return filter_nan( + torch.arange(0, 2**16, dtype=torch.int32, device="cuda") + .to(torch.uint16) + .view(dtype) + ) + if dtype == torch.float32: + if slow: + return filter_nan( + torch.arange(0, 2**32, dtype=torch.int64, device="cuda") + .to(torch.uint32) + .view(dtype) + ) + else: + # E8M5, E5M10, E8M23 + result = [] + for t in [ + generate_test_tensor(torch.bfloat16), + generate_test_tensor(torch.float16), + ]: + orig_dtype = t.dtype + t = t.to(dtype).view(torch.uint32) + missing_mantissa_bits = mantissa_bits(dtype) - mantissa_bits(orig_dtype) + r = torch.randint( + 0, 2**missing_mantissa_bits, t.shape, dtype=t.dtype, device="cuda" + ) + result.append(t.view(dtype)) + result.append((t.view(torch.int32) | r.view(torch.int32)).view(dtype)) + return filter_nan(torch.cat(result)) + + +def make_input(dtype, category): + dtype = getattr(torch, dtype) + from torch.testing import make_tensor + + if category in ["matrix"]: + return make_tensor(1024, 1024, dtype=dtype, device="cuda"), make_tensor( + 1024, 1024, dtype=dtype, device="cuda" + ) + if category == "normalization": + return make_tensor(1024, 1024, dtype=dtype, device="cuda"), make_tensor( + 1024, dtype=dtype, device="cuda" + ) + if category == "reduction": + return (make_tensor(1024, 1024, dtype=dtype, device="cuda"),) + if category == "binary": + test_tensor = generate_test_tensor(dtype) + # randomly permute twice: + return ( + test_tensor[torch.randperm(test_tensor.shape[0])], + test_tensor[torch.randperm(test_tensor.shape[0])], + ) + return (generate_test_tensor(dtype),) + + +def make_function(function, category): + if category == "reduction": + return eval(f"lambda x: {function}(x, dim=-1)") + elif category == "matrix": + return eval(f"lambda x, y: {function}(x, y)") + elif category == "normalization": + return eval(f"lambda x, y: {function}(x, [x.shape[-1]], y, eps=1e-5)") + elif category == "binary": + return eval(f"lambda x, y: {function}(x, y)") + elif category == "activation": + return eval(f"lambda x: {function}(x)") + elif category == "elementary": + return eval(f"lambda x: {function}(x)") + else: + raise ValueError(f"Unknown category: {category}") + + +def rel_err_ulp(a, b, dtype): + return (a - b).abs() / (b.abs() * torch.finfo(dtype).eps + torch.finfo(dtype).tiny) + + +def evaluate_output(input, output, golden_output): + # todo: how to make sure this handles denormals in the input well? + # for now, only handle it if the inputs have the same shape, otherwise assume there are none + # handle more than one input + + input = [i.flatten().float() for i in input if i.shape == output.shape] + + dtype = golden_output.dtype + + output = output.flatten().float() + golden_output = golden_output.flatten().float() + + # we are checking subnormals separate from the rest of the numbers + # we also need to check NaNs and Infs carefully + + subnormal_mask = golden_output.abs() < torch.finfo(dtype).smallest_normal + for i in input: + subnormal_mask |= i.abs() < torch.finfo(dtype).smallest_normal + nan_mask_golden = torch.isnan(golden_output) + golden_output = torch.where(nan_mask_golden, float("nan"), golden_output) + nan_mask_output = torch.isnan(output) + output = torch.where(nan_mask_output, float("nan"), output) + + equal_mask = (output == golden_output) | (nan_mask_golden & nan_mask_output) + equal_subnormal_mask = ( + torch.where(subnormal_mask, 0, output) + == torch.where(subnormal_mask, 0, golden_output) + ) | equal_mask + output_flushed = torch.where(subnormal_mask, 0.0, output) + golden_flushed = torch.where(subnormal_mask, 0.0, golden_output) + output_flushed = torch.where(~torch.isfinite(output_flushed), 0.0, output_flushed) + golden_flushed = torch.where(~torch.isfinite(golden_flushed), 0.0, golden_flushed) + + num_nonequal = (~equal_mask).sum().item() + num_nonequal_subnormal = (~equal_subnormal_mask).sum().item() + + err = rel_err_ulp(output_flushed, golden_flushed, dtype) + max_ulp_to_golden = err.max().item() + avg_ulp_to_golden = err.mean().item() + mismatch_sample = [] + pos = (~equal_mask).nonzero().squeeze(1) + log.info( + "pos.shape: %s, input: %s, golden_output.shape: %s, output.shape: %s", + pos.shape, + [i.shape for i in input], + golden_output.shape, + output.shape, + ) + ordered = torch.argsort(err[pos], descending=True) + random = torch.randperm(pos.shape[0]) + for sampling in [ordered, random]: + for i in range(min(pos.shape[0], 5)): + sample_idx = pos[sampling[i]] + mismatch_sample.append( + { + "pos": sample_idx.item(), + "input": [inp[sample_idx].item() for inp in input], + "output": output[sample_idx].item(), + "golden": golden_output[sample_idx].item(), + "rel_err": err[sample_idx].item(), + } + ) + + # use hashlib.sha256 to hash the tensors + # this i + return { + "normal_hash": hashlib.sha256( + output_flushed.cpu().numpy().tobytes() + ).hexdigest()[:8], + "full_hash": hashlib.sha256(output.cpu().numpy().tobytes()).hexdigest()[:8], + "max_ulp_to_golden": max_ulp_to_golden, + "avg_ulp_to_golden": avg_ulp_to_golden, + "num_nonequal": num_nonequal, + "num_nonequal_subnormal": num_nonequal_subnormal, + "num_total": output.shape[0], + "match_full": num_nonequal == 0, + "match_normal": num_nonequal_subnormal == 0, + "mismatch_sample": mismatch_sample, + } + + +def create_golden(args): + golden_inputs = { + dtype: {category: make_input(dtype, category) for category in CATEGORIES} + for dtype in DTYPES + } + golden_outputs = { + dtype: { + category: { + function: { + pass_type: MODES[args.mode]( + PASSES[pass_type](make_function(function, category)) + )(*golden_inputs[dtype][category]) + for pass_type in PASS_TYPES[category] + } + for function in CATEGORIES[category] + } + for category in CATEGORIES + } + for dtype in DTYPES + } + return golden_inputs, golden_outputs + + +def get_metadata(args, dtype, category, function, pass_type): + metadata = { + "gpu": args.gpu[0], + "pytorch_version": args.pytorch_version, + "cuda_version": args.cuda_version, + "pytorch_mode": args.mode, + "is_golden": args.create_golden, + "data_type": dtype, + "function": function, + "pass_type": pass_type, + "category": category, + } + identifier = hashlib.sha256(json.dumps(metadata).encode()).hexdigest()[:8] + metadata["identifier"] = identifier + return metadata + + +def run_test_case(args): + golden = torch.load(args.golden) + golden_inputs, golden_outputs = golden + metadata = get_metadata( + args, args.dtype, args.category, args.function, args.pass_type + ) + input = golden_inputs[args.dtype][args.category] + callable = make_function(args.function, args.category) + callable = PASSES[args.pass_type](callable) + callable = MODES[args.mode](callable) + output = callable(*input) + evaluation = evaluate_output( + input, + output, + golden_outputs[args.dtype][args.category][args.function][args.pass_type], + ) + data = json.dumps(metadata | evaluation) + log.info(data) + print(data) + + +test_concurrency = asyncio.Semaphore(16) + + +async def launch_test_case(args, dtype, category, function, pass_type): + async with test_concurrency: + env = os.environ | MODES[args.mode].env + metadata = get_metadata(args, dtype, category, function, pass_type) + identifier = metadata["identifier"] + env["TORCH_TRACE"] = f"/workspace/logs/{identifier}.trace" + Path(env["TORCH_TRACE"]).mkdir(parents=True, exist_ok=True) + env["TORCH_LOGS"] = "+all" + env["TORCHINDUCTOR_FORCE_DISABLE_CACHES"] = "1" + env["TRITON_ALWAYS_COMPILE"] = "1" + cmd_args = [ + "--test", + f"--gpu {args.gpu[0]}", + f"--pytorch-version {args.pytorch_version}", + f"--cuda-version {args.cuda_version}", + f"--mode {args.mode}", + f"--golden {args.golden}", + f"--dtype {dtype}", + f"--category {category}", + f"--function {function}", + f"--pass-type {pass_type}", + ] + if args.create_golden: + cmd_args.append("--create-golden") + return await run( + f"{sys.executable} {__file__} {' '.join(cmd_args)}", + env=env, + capture_stderr=False, + stderr=open(f"/workspace/logs/{identifier}.log", "w"), + ) + + +async def runner(args): + log.info("Runner called! %s", args) + + if args.create_golden: + golden = create_golden(args) + torch.save(golden, args.golden) + + results = await asyncio.gather( + *[ + launch_test_case( + args, + dtype, + category, + function, + pass_type, + ) + for dtype, category in itertools.product(DTYPES, CATEGORIES) + for function, pass_type in itertools.product( + CATEGORIES[category], PASS_TYPES[category] + ) + ] + ) + + with open( + f"/workspace/result_{args.gpu[0]}_{args.pytorch_version}_{args.cuda_version}.jsonl", + "a+", + ) as result: + for r in results: + result.write(r) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--launcher", + default=False, + action="store_true", + help="launch workers across machines", + ) + parser.add_argument( + "--worker", + default=False, + action="store_true", + help="launches jobs on a single machine", + ) + parser.add_argument( + "--runner", + default=False, + action="store_true", + help="runs tests in a specific environment", + ) + parser.add_argument( + "--test", + action="store_true", + help="runs a single test case", + ) + parser.add_argument("--gpu", nargs="+", type=str, default=["t4"]) + parser.add_argument("--html", type=str, default=None) + parser.add_argument("--golden", type=str, default=None) + parser.add_argument("--create-golden", action="store_true") + parser.add_argument("--dtype", type=str) + parser.add_argument("--category", type=str) + parser.add_argument("--function", type=str) + parser.add_argument("--pass-type", type=str) + parser.add_argument("--pytorch-version", type=str) + parser.add_argument("--cuda-version", type=str) + parser.add_argument("--mode") + args = parser.parse_args() + if args.runner: + assert torch is not None # noqa: S101 + torch.set_default_device("cuda") + (gpu,) = args.gpu + logging.basicConfig( + level=logging.INFO, + format=f"%(asctime)s - runner:{gpu}/{args.pytorch_version}/{args.cuda_version}/{args.mode} - %(message)s", # noqa: B950 + ) + asyncio.run(runner(args)) + if args.worker: + (gpu,) = args.gpu + logging.basicConfig( + level=logging.INFO, format=f"%(asctime)s - worker:{gpu} - %(message)s" + ) + asyncio.run(worker(args)) + if args.launcher: + logging.basicConfig( + level=logging.INFO, format="%(asctime)s - launcher - %(message)s" + ) + try: + asyncio.run(launcher(args)) + except KeyboardInterrupt: + log.error("Cancelled by user") + if args.test: + test_id = ( + f"{args.gpu[0]}/{args.pytorch_version}/{args.cuda_version}" + f"/{args.mode}/{args.dtype}/{args.category}" + f"/{args.function}/{args.pass_type}" + ) + logging.basicConfig( + level=logging.INFO, + format=f"%(asctime)s - test:{test_id} - %(message)s", + ) + run_test_case(args) + if args.html: + format_results_to_html(args.html) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/dynamo/perf_cli.py b/benchmarks/dynamo/perf_cli.py new file mode 100644 index 0000000000000..f4c85a5e4a2a8 --- /dev/null +++ b/benchmarks/dynamo/perf_cli.py @@ -0,0 +1,1894 @@ +#!/usr/bin/env python3 +""" +CLI tool for launching, summarizing, and reproducing inductor perf regression runs. + +Usage: + python benchmarks/dynamo/perf_cli.py launch [--device a100 h100 ...] [--ref BRANCH] [--wait] + python benchmarks/dynamo/perf_cli.py summary [--top 5] [--config PATTERN] + python benchmarks/dynamo/perf_cli.py repro [--model MODEL] [--suite SUITE] [--print-only] + +Requires: gh CLI (authenticated), internet access to S3 (gha-artifacts bucket). +""" + +from __future__ import annotations + +import argparse +import csv +import io +import json +import os +import re +import shutil +import subprocess +import sys +import time +import urllib.request +import zipfile +from collections import defaultdict +from dataclasses import dataclass, field +from datetime import datetime +from math import exp, log +from pathlib import Path + + +WORKFLOWS = { + "a100": { + "name": "inductor-A100-perf-nightly", + "id": 42513231, + }, + "a100-compare": { + "name": "inductor-A100-perf-compare", + "id": 50531883, + }, + "h100": { + "name": "inductor-perf-nightly-h100", + "id": 144201955, + }, + "b200": { + "name": "inductor-perf-b200", + "id": 173716622, + }, + "rocm-mi300": { + "name": "inductor-perf-nightly-rocm-mi300", + "id": 197925166, + }, + "rocm-mi355": { + "name": "inductor-perf-nightly-rocm-mi355", + "id": 197925165, + }, + "x86": { + "name": "inductor-perf-nightly-x86", + "id": 108782874, + }, + "x86-zen": { + "name": "inductor-perf-nightly-x86-zen", + "id": 167573808, + }, + "aarch64": { + "name": "inductor-perf-nightly-aarch64", + "id": 109196799, + }, + "macos": { + "name": "inductor-perf-nightly-macos", + "id": 117199085, + }, + "xpu": { + "name": "inductor-perf-nightly-xpu", + "id": 201149053, + }, +} + +DEVICE_CHOICES = sorted(k for k in WORKFLOWS if k != "a100-compare") + +S3_BUCKET = "gha-artifacts" +S3_URL = f"https://{S3_BUCKET}.s3.amazonaws.com" +REPO = "pytorch/pytorch" + +# Regex to parse test job names like: +# "cuda13.0-py3.10-gcc11-sm80 / test (inductor_huggingface_perf, 1, 5, linux.aws.a100)" +JOB_RE = re.compile( + r"test \((?P[^,]+),\s*(?P\d+),\s*(?P\d+),\s*(?P[^)]+)\)" +) + +PERF_CONFIGS = re.compile(r"inductor_(huggingface|timm|torchbench)_perf") + +SUITE_ALIASES = { + "hf": "huggingface", + "huggingface": "huggingface", + "timm": "timm_models", + "timm_models": "timm_models", + "tb": "torchbench", + "torchbench": "torchbench", +} + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def gh(*args: str, json_output: bool = False) -> str | dict | list: + cmd = ["gh"] + list(args) + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + print(f"gh error: {result.stderr.strip()}", file=sys.stderr) + sys.exit(1) + if json_output: + return json.loads(result.stdout) + return result.stdout.strip() + + +def git(*args: str) -> str: + result = subprocess.run(["git"] + list(args), capture_output=True, text=True) + return result.stdout.strip() + + +def gmean(values: list[float]) -> float: + if not values: + return 0.0 + return exp(sum(log(v) for v in values if v > 0) / max(len(values), 1)) + + +@dataclass +class Metric: + name: str + field: str # attribute on ModelResult + unit: str # display suffix + higher_is_better: bool + aggregate: str # "gmean" or "mean" + + +METRICS = { + "speedup": Metric("speedup", "speedup", "x", True, "gmean"), + "compilation_latency": Metric( + "compilation latency", "compilation_latency", "s", False, "mean" + ), + "compression_ratio": Metric( + "memory compression", "compression_ratio", "x", True, "gmean" + ), + "abs_latency": Metric("absolute latency", "abs_latency", "ms", False, "mean"), +} + +METRIC_CHOICES = list(METRICS.keys()) + +# HUD uses 5% relative threshold for flagging regressions +RELATIVE_THRESHOLD = 0.05 + +WORKFLOW_NAME_TO_NIGHTLY_ID = { + v["name"]: v["id"] for k, v in WORKFLOWS.items() if k != "a100-compare" +} +# compare's baseline is A100 nightly +WORKFLOW_NAME_TO_NIGHTLY_ID["inductor-A100-perf-compare"] = WORKFLOWS["a100"]["id"] + + +def _short_config(config: str, device: str = "") -> str: + """Compact label: e.g. 'a100 cudagraphs huggingface training'.""" + c = config + backend = c + for s in ("_huggingface_", "_timm_models_", "_torchbench_"): + if s in c: + backend = c.split(s)[0] + break + backend = backend.removeprefix("inductor_") + mode = "training" if "training" in c else "inference" if "inference" in c else "" + suite = "" + for s in ("huggingface", "timm_models", "torchbench"): + if s in c: + suite = s + break + parts = [p for p in (device, backend, suite, mode) if p] + return " ".join(parts) + + +@dataclass +class ModelResult: + name: str + speedup: float + abs_latency: float = 0.0 + compilation_latency: float = 0.0 + compression_ratio: float = 0.0 + eager_peak_mem: float = 0.0 + dynamo_peak_mem: float = 0.0 + config: str = "" + device: str = "" + + @property + def short_config(self) -> str: + return _short_config(self.config, self.device) + + +@dataclass +class PerfData: + config: str # e.g. "inductor_with_cudagraphs_huggingface_amp_training_cuda" + models: list[ModelResult] = field(default_factory=list) + device: str = "" + + @property + def suite(self) -> str: + for s in ("huggingface", "timm_models", "torchbench"): + if s in self.config: + return s + return "unknown" + + @property + def mode(self) -> str: + if "training" in self.config: + return "training" + if "inference" in self.config: + return "inference" + return "unknown" + + @property + def dtype(self) -> str: + # Config format: {backend}_{suite}_{dtype}_{mode}_{device} + # e.g. inductor_with_cudagraphs_huggingface_amp_training_cuda + for s in ("_huggingface_", "_timm_models_", "_torchbench_"): + if s in self.config: + tail = self.config.split(s, 1)[1] + # tail is e.g. "amp_training_cuda" + parts = tail.split("_") + if parts: + return parts[0] + return "unknown" + + @property + def short_name(self) -> str: + name = self.config + for s in ("_huggingface_", "_timm_models_", "_torchbench_"): + if s in name: + name = name.split(s)[0] + break + return name + + @property + def runtime(self) -> str: + # Last token in config: e.g. "..._training_cuda" → "cuda" + parts = self.config.rsplit("_", 1) + if len(parts) == 2: + return parts[1] + return "unknown" + + @property + def qualified_config(self) -> str: + if self.device: + return f"{self.device}/{self.config}" + return self.config + + def gmean_speedup(self) -> float: + vals = [m.speedup for m in self.models if m.speedup > 0] + return gmean(vals) + + def aggregate_metric(self, metric: Metric) -> float: + vals = [ + getattr(m, metric.field) + for m in self.models + if getattr(m, metric.field) > 0 + ] + if not vals: + return 0.0 + if metric.aggregate == "gmean": + return gmean(vals) + return sum(vals) / len(vals) + + +# --------------------------------------------------------------------------- +# Artifact downloading +# --------------------------------------------------------------------------- + + +def get_run_jobs(run_id: int) -> list[dict]: + data = gh( + "run", + "view", + str(run_id), + "--repo", + REPO, + "--json", + "jobs", + json_output=True, + ) + return data["jobs"] + + +def get_perf_jobs(jobs: list[dict]) -> list[dict]: + perf_jobs = [] + for job in jobs: + m = JOB_RE.search(job["name"]) + if not m: + continue + config = m.group("config") + if not PERF_CONFIGS.match(config): + continue + if job.get("conclusion") != "success": + continue + perf_jobs.append( + { + "config": config, + "shard": m.group("shard"), + "num_shards": m.group("num_shards"), + "runner": m.group("runner"), + "job_id": job["databaseId"], + "name": job["name"], + } + ) + return perf_jobs + + +def s3_artifact_url(run_id: int, attempt: int, job: dict) -> str: + config = job["config"] + shard = job["shard"] + num_shards = job["num_shards"] + runner = job["runner"] + job_id = job["job_id"] + filename = f"test-reports-test-{config}-{shard}-{num_shards}-{runner}_{job_id}.zip" + return f"{S3_URL}/{REPO}/{run_id}/{attempt}/artifact/{filename}" + + +CACHE_DIR = Path.home() / ".cache" / "perf_cli" + + +def get_cache_dir(run_id: int, attempt: int) -> Path: + d = CACHE_DIR / f"{run_id}" / f"{attempt}" + d.mkdir(parents=True, exist_ok=True) + return d + + +def download_and_extract_csvs( + run_id: int, + jobs: list[dict], + attempt: int = 1, + no_cache: bool = False, +) -> list[tuple[str, str]]: + """Download artifacts and return list of (csv_filename, csv_content) pairs.""" + cache = get_cache_dir(run_id, attempt) + results = [] + fetched = 0 + + for job in jobs: + # Check cache first + cache_key = ( + f"{job['config']}-{job['shard']}-{job['num_shards']}-{job['job_id']}" + ) + cache_marker = cache / f"{cache_key}.done" + + if not no_cache and cache_marker.exists(): + # Read cached CSVs + for csv_file in cache.glob(f"{cache_key}__*.csv"): + csv_name = csv_file.name.split("__", 1)[1] + results.append((csv_name, csv_file.read_text())) + continue + + url = s3_artifact_url(run_id, attempt, job) + zip_path = cache / f"{cache_key}.zip" + try: + urllib.request.urlretrieve(url, str(zip_path)) + fetched += 1 + except urllib.error.HTTPError as e: + print( + f" warning: failed to download shard {job['config']} " + f"shard {job['shard']}: {e}", + file=sys.stderr, + ) + continue + + with zipfile.ZipFile(zip_path) as zf: + for name in zf.namelist(): + if name.endswith("_performance.csv"): + csv_name = os.path.basename(name) + with zf.open(name) as f: + content = f.read().decode("utf-8") + results.append((csv_name, content)) + # Write to cache + (cache / f"{cache_key}__{csv_name}").write_text(content) + + # Mark this shard as cached and remove the zip + cache_marker.touch() + zip_path.unlink(missing_ok=True) + + if fetched > 0: + print(f" downloaded {fetched} shards (cached at {cache})") + elif results: + print(f" using cached data from {cache}") + + return results + + +def parse_csvs(csv_pairs: list[tuple[str, str]], device: str = "") -> list[PerfData]: + grouped: dict[str, list[ModelResult]] = defaultdict(list) + + for csv_name, content in csv_pairs: + # csv_name like: inductor_with_cudagraphs_huggingface_amp_training_cuda_performance.csv + config = csv_name.replace("_performance.csv", "") + reader = csv.DictReader(io.StringIO(content)) + for row in reader: + try: + speedup = float(row.get("speedup", 0)) + except (ValueError, TypeError): + continue + grouped[config].append( + ModelResult( + name=row.get("name", "?"), + speedup=speedup, + abs_latency=float(row.get("abs_latency", 0) or 0), + compilation_latency=float(row.get("compilation_latency", 0) or 0), + compression_ratio=float(row.get("compression_ratio", 0) or 0), + eager_peak_mem=float(row.get("eager_peak_mem", 0) or 0), + dynamo_peak_mem=float(row.get("dynamo_peak_mem", 0) or 0), + config=config, + device=device, + ) + ) + + return [ + PerfData(config=k, models=v, device=device) for k, v in sorted(grouped.items()) + ] + + +# --------------------------------------------------------------------------- +# S-curve rendering +# --------------------------------------------------------------------------- + + +def subsample(items: list, max_rows: int) -> list: + """Evenly subsample a sorted list, always keeping first and last.""" + n = len(items) + if n <= max_rows: + return items + # Always include first and last; evenly space the rest + indices = {0, n - 1} + for i in range(1, max_rows - 1): + indices.add(round(i * (n - 1) / (max_rows - 1))) + return [items[i] for i in sorted(indices)] + + +def render_scurve( + perf: PerfData, + metric: Metric, + top_n: int = 5, + term_width: int | None = None, + term_height: int | None = None, +): + if not perf.models: + return + + if term_width is None or term_height is None: + sz = shutil.get_terminal_size((100, 50)) + term_width = term_width or sz.columns + term_height = term_height or sz.lines + + live = [m for m in perf.models if getattr(m, metric.field) > 0] + if not live: + return + + sorted_models = sorted(live, key=lambda m: getattr(m, metric.field)) + agg = perf.aggregate_metric(metric) + n = len(sorted_models) + + # Reserve lines for header (2) + axis label (1) + padding (2) + max_rows = max(term_height - 5, 15) + display = subsample(sorted_models, max_rows) + skipped = n - len(display) + + agg_label = metric.aggregate + header = f"{perf.config} ({n} data points, {agg_label}={agg:.2f}{metric.unit})" + if skipped > 0: + header += f" [showing {len(display)}/{n}]" + print(f"\n {header}") + print(f" {'─' * min(len(header), term_width - 4)}") + + def fmt_val(v: float) -> str: + if metric.unit == "s" or metric.unit == "ms": + return f"{v:7.1f}{metric.unit}" + return f"{v:5.2f}{metric.unit}" + + # Layout: " {name:8} {dots}" + sample_val = fmt_val(display[0] and getattr(display[0], metric.field)) + max_name = min(max(len(m.name) for m in display), 30) + val_width = len(sample_val) + prefix_len = 2 + max_name + 2 + val_width + 2 + plot_width = max(term_width - prefix_len - 1, 20) + + def get_val(m): + return getattr(m, metric.field) + + min_val = get_val(sorted_models[0]) + p95_idx = max(0, int(n * 0.95) - 1) + p95_val = get_val(sorted_models[p95_idx]) + + # For ratio metrics (speedup, compression_ratio), anchor at 1.0 + # For absolute metrics (latency), anchor at 0 + if metric.unit == "x": + plot_min = min(min_val, 0.5) + plot_max = max(p95_val * 1.1, 1.5) + marker_val = 1.0 + marker_label = "1.0x" + else: + plot_min = 0 + plot_max = p95_val * 1.1 + marker_val = None + marker_label = None + + span = plot_max - plot_min + if span == 0: + span = 1 + + def val_to_col(v: float) -> int: + return max( + 0, min(plot_width - 1, int((v - plot_min) / span * (plot_width - 1))) + ) + + marker_col = val_to_col(marker_val) if marker_val is not None else None + + for m in display: + name = m.name[:max_name].ljust(max_name) + v = get_val(m) + col = val_to_col(v) + bar = [" "] * plot_width + if marker_col is not None: + bar[marker_col] = "|" + for i in range(col + 1): + if marker_col is not None and i == marker_col: + bar[i] = "|" + else: + bar[i] = "·" + print(f" {name} {fmt_val(v)} {''.join(bar)}") + + pad = " " * prefix_len + if marker_label and marker_col is not None: + print(f"{pad}{' ' * marker_col}{marker_label}") + else: + print() + + +def print_worst_offenders(perf: PerfData, metric: Metric, top_n: int = 5): + def get_val(m): + return getattr(m, metric.field) + + live = [m for m in perf.models if get_val(m) > 0] + if not live: + return + # "worst" depends on metric direction + if metric.higher_is_better: + worst = sorted(live, key=get_val)[:top_n] + else: + worst = sorted(live, key=get_val, reverse=True)[:top_n] + print(f"\n Worst offenders ({metric.name}):") + for i, m in enumerate(worst, 1): + v = get_val(m) + parts = [f"{v:.3f}{metric.unit}"] + if m.config: + parts.append(m.short_config) + detail = " ".join(parts) + print(f" {i}. {m.name:<30} {detail}") + + +# --------------------------------------------------------------------------- +# Subcommands +# --------------------------------------------------------------------------- + + +def build_dispatch_inputs(args) -> list[str]: + """Build -f flags for workflow dispatch inputs from CLI args.""" + flags = [] + bool_inputs = [ + "training", + "inference", + "default", + "dynamic", + "cppwrapper", + "cudagraphs", + "freezing_cudagraphs", + "aotinductor", + "maxautotune", + ] + for name in bool_inputs: + val = getattr(args, name, None) + if val is not None: + flags.extend(["-f", f"{name}={'true' if val else 'false'}"]) + if args.benchmark_configs: + flags.extend(["-f", f"benchmark_configs={args.benchmark_configs}"]) + return flags + + +def dispatch_one(device: str, ref: str, extra_flags: list[str]) -> int | None: + wf = WORKFLOWS[device] + print(f"\nLaunching {wf['name']} on ref: {ref}") + + dispatch_args = [ + "workflow", + "run", + str(wf["id"]), + "--repo", + REPO, + "--ref", + ref, + ] + extra_flags + + gh(*dispatch_args) + print("Dispatched. Waiting a few seconds for the run to appear...") + time.sleep(5) + + runs = gh( + "run", + "list", + "--repo", + REPO, + "--workflow", + str(wf["id"]), + "--branch", + ref, + "--limit", + "1", + "--json", + "databaseId,status,url,createdAt", + json_output=True, + ) + if not runs: + print("Could not find the dispatched run. Check the Actions tab manually.") + return None + + run = runs[0] + run_id = run["databaseId"] + url = f"https://github.com/{REPO}/actions/runs/{run_id}" + print(f"Run ID: {run_id}") + print(f"URL: {url}") + print(f"Status: {run['status']}") + return run_id + + +def wait_for_runs(pending: dict[str, int]) -> dict[str, int]: + """Poll all runs until they complete. Returns dict of successful runs.""" + print(f"\nWaiting for {len(pending)} run(s)...", flush=True) + remaining = dict(pending) + succeeded: dict[str, int] = {} + while remaining: + time.sleep(30) + done = [] + for device, run_id in remaining.items(): + data = gh( + "run", + "view", + str(run_id), + "--repo", + REPO, + "--json", + "status,conclusion", + json_output=True, + ) + status = data.get("status", "unknown") + if status == "completed": + conclusion = data.get("conclusion", "unknown") + print(f" {device} (run {run_id}): {conclusion}") + if conclusion == "success": + succeeded[device] = run_id + done.append(device) + for d in done: + del remaining[d] + if remaining: + ts = datetime.now().strftime("%H:%M:%S") + print(f" [{ts}] still waiting: {', '.join(remaining)}...", flush=True) + return succeeded + + +def cmd_launch(args): + ref = args.ref or git("rev-parse", "--abbrev-ref", "HEAD") + if ref == "HEAD": + ref = git("rev-parse", "HEAD") + + extra_flags = build_dispatch_inputs(args) + launched: dict[str, int] = {} + for device in args.device: + if device not in WORKFLOWS: + print(f"Unknown device: {device}", file=sys.stderr) + sys.exit(1) + run_id = dispatch_one(device, ref, extra_flags) + if run_id: + launched[device] = run_id + + if (args.wait or args.wait_and_summarize) and launched: + succeeded = wait_for_runs(launched) + if args.wait_and_summarize and succeeded: + # Use the first device's run ID as the positional arg; pass all + # device→run_id pairs via _run_ids so cmd_summary skips resolution. + first_run = next(iter(succeeded.values())) + summary_args = argparse.Namespace( + run_id=str(first_run), + device=list(succeeded.keys()), + _run_ids=succeeded, + baseline="latest", + metric="speedup", + top=5, + config=None, + suite=None, + mode=None, + group_by=None, + attempt=1, + no_cache=False, + ) + print(f"\n{'=' * 70}") + print("Summary") + print(f"{'=' * 70}") + cmd_summary(summary_args) + + +def filter_perf(all_perf: list[PerfData], args) -> list[PerfData]: + result = all_perf + if getattr(args, "config", None): + pattern = re.compile(args.config, re.IGNORECASE) + result = [p for p in result if pattern.search(p.config)] + if getattr(args, "suite", None): + suite = SUITE_ALIASES.get(args.suite, args.suite) + result = [p for p in result if p.suite == suite] + if getattr(args, "mode", None): + result = [p for p in result if p.mode == args.mode] + if getattr(args, "backend", None): + pattern = re.compile(args.backend, re.IGNORECASE) + result = [p for p in result if pattern.search(p.short_name)] + if getattr(args, "dtype", None): + result = [p for p in result if p.dtype == args.dtype] + if getattr(args, "runtime", None): + result = [p for p in result if p.runtime == args.runtime] + return result + + +GROUP_KEY_FNS: dict[str, callable] = { + "config": lambda p: p.qualified_config, + "suite": lambda p: p.suite, + "mode": lambda p: p.mode, + "backend": lambda p: p.short_name, + "device": lambda p: p.device or "unknown", + "dtype": lambda p: p.dtype, + "runtime": lambda p: p.runtime, +} + +GROUP_CHOICES = sorted(GROUP_KEY_FNS.keys()) + + +def group_perf(all_perf: list[PerfData], group_by: list[str] | None) -> list[PerfData]: + if not group_by: + all_models = [] + for p in all_perf: + all_models.extend(p.models) + return [PerfData(config="all", models=all_models)] + + fns = [] + for key in group_by: + fn = GROUP_KEY_FNS.get(key) + if fn is None: + print(f"Unknown group-by: {key}", file=sys.stderr) + sys.exit(1) + fns.append(fn) + + def composite_key(p: PerfData) -> str: + return " / ".join(fn(p) for fn in fns) + + groups: dict[str, list[ModelResult]] = defaultdict(list) + for p in all_perf: + groups[composite_key(p)].extend(p.models) + + return [PerfData(config=k, models=v) for k, v in sorted(groups.items())] + + +# --------------------------------------------------------------------------- +# Baseline comparison +# --------------------------------------------------------------------------- + + +NIGHTLY_WORKFLOW_IDS = {k: v["id"] for k, v in WORKFLOWS.items() if k != "a100-compare"} + + +def _find_latest_run( + branch: str, + device: str, +) -> dict | None: + """Find the latest successful perf nightly run for a branch + device. + + Returns {databaseId, createdAt, headSha} or None. + """ + wf_id = NIGHTLY_WORKFLOW_IDS.get(device) + if wf_id is None: + return None + + runs = gh( + "run", + "list", + "--repo", + REPO, + "--workflow", + str(wf_id), + "--branch", + branch, + "--status", + "success", + "--limit", + "1", + "--json", + "databaseId,createdAt,headSha", + json_output=True, + ) + if not runs: + return None + return runs[0] + + +def resolve_run(branch: str, device: str) -> int: + """Find the latest successful perf nightly run for a branch + device.""" + run = _find_latest_run(branch, device) + if not run: + print( + f"No successful perf run found for branch '{branch}' on {device}.", + file=sys.stderr, + ) + sys.exit(1) + run_id = run["databaseId"] + created = run["createdAt"][:10] + print(f"Resolved '{branch}' → run {run_id} ({device}, {created})") + return run_id + + +def discover_runs(branch: str) -> dict[str, int]: + """Auto-discover all devices with successful runs for a branch. + + Finds the latest commit SHA that has runs, then returns all runs matching + that commit. + """ + candidates: list[tuple[str, dict]] = [] + for device in DEVICE_CHOICES: + run = _find_latest_run(branch, device) + if run: + candidates.append((device, run)) + + if not candidates: + print( + f"No successful perf runs found for branch '{branch}' on any device.", + file=sys.stderr, + ) + sys.exit(1) + + # Pick the most recent commit (by createdAt) and collect all runs on it + candidates.sort(key=lambda x: x[1]["createdAt"], reverse=True) + target_sha = candidates[0][1]["headSha"] + + result = {} + for device, run in candidates: + if run["headSha"] == target_sha: + run_id = run["databaseId"] + created = run["createdAt"][:10] + print( + f"Discovered '{branch}' → run {run_id} ({device}, {created}, {target_sha[:10]})" + ) + result[device] = run_id + + return result + + +def resolve_runs(branch: str, devices: list[str]) -> dict[str, int]: + """Resolve the latest successful run for each device. Returns {device: run_id}.""" + result = {} + for device in devices: + result[device] = resolve_run(branch, device) + return result + + +def resolve_baseline(head_run_id: int) -> int: + """Find the latest successful nightly on main, skipping the head run itself.""" + run_data = gh( + "run", + "view", + str(head_run_id), + "--repo", + REPO, + "--json", + "workflowName", + json_output=True, + ) + wf_name = run_data.get("workflowName", "") + nightly_id = WORKFLOW_NAME_TO_NIGHTLY_ID.get(wf_name) + if nightly_id is None: + print( + f"Don't know which nightly corresponds to workflow '{wf_name}'", + file=sys.stderr, + ) + sys.exit(1) + + runs = gh( + "run", + "list", + "--repo", + REPO, + "--workflow", + str(nightly_id), + "--branch", + "main", + "--status", + "success", + "--limit", + "5", + "--json", + "databaseId,createdAt,headBranch", + json_output=True, + ) + # Skip the head run itself to avoid comparing a run to itself + for run in runs: + if run["databaseId"] != head_run_id: + baseline_id = run["databaseId"] + created = run["createdAt"][:10] + print(f"Baseline: run {baseline_id} (main, {created})") + return baseline_id + + print("No suitable baseline nightly found on main.", file=sys.stderr) + sys.exit(1) + + +def device_for_workflow(workflow_name: str) -> str: + """Reverse-lookup device key from workflow name.""" + for k, v in WORKFLOWS.items(): + if v["name"] == workflow_name: + return k + return "" + + +@dataclass +class RunMeta: + run_id: int + head_sha: str + head_branch: str + workflow_name: str + created_at: str + event: str + + @property + def short_sha(self) -> str: + return self.head_sha[:10] + + @property + def date(self) -> str: + return self.created_at[:10] + + +def fetch_run_meta(run_id: int) -> RunMeta: + data = gh( + "api", + f"repos/{REPO}/actions/runs/{run_id}", + "-q", + "{headSha: .head_sha, headBranch: .head_branch, workflowName: .name, createdAt: .created_at, event: .event}", + json_output=True, + ) + return RunMeta( + run_id=run_id, + head_sha=data.get("headSha", "unknown"), + head_branch=data.get("headBranch", "unknown"), + workflow_name=data.get("workflowName", "unknown"), + created_at=data.get("createdAt", "unknown"), + event=data.get("event", "unknown"), + ) + + +def print_run_header( + label: str, + metas: list[RunMeta], + configs: list[str] | None = None, +): + print(f"\n {label}") + print(f" {'─' * len(label)}") + if len(metas) == 1: + m = metas[0] + print(f" Run: {m.run_id} ({m.workflow_name})") + print(f" Commit: {m.short_sha} ({m.head_branch}, {m.date})") + else: + # Show commit from first (should all match for multi-device) + print( + f" Commit: {metas[0].short_sha} ({metas[0].head_branch}, {metas[0].date})" + ) + print(" Runs:") + for m in metas: + print(f" {m.run_id} ({m.workflow_name})") + if configs: + print(f" Configs: {len(configs)} — {', '.join(sorted(configs))}") + + +def fetch_run_perf( + run_id: int, + attempt: int, + no_cache: bool, + device: str = "", + allow_empty: bool = False, +) -> list[PerfData]: + """Fetch and parse perf data for a run.""" + jobs = get_run_jobs(run_id) + perf_jobs = get_perf_jobs(jobs) + if not perf_jobs: + if allow_empty: + print(f" no perf jobs in run {run_id}, skipping", file=sys.stderr) + return [] + print(f"No successful perf jobs in run {run_id}.", file=sys.stderr) + sys.exit(1) + csv_pairs = download_and_extract_csvs(run_id, perf_jobs, attempt, no_cache=no_cache) + if not csv_pairs: + if allow_empty: + print(f" no CSVs in run {run_id}, skipping", file=sys.stderr) + return [] + print(f"No performance CSVs in run {run_id}.", file=sys.stderr) + sys.exit(1) + return parse_csvs(csv_pairs, device=device) + + +@dataclass +class ModelDelta: + name: str + base_val: float + head_val: float + config: str = "" + device: str = "" + + @property + def delta_pct(self) -> float: + if self.base_val == 0: + return 0.0 + return (self.head_val - self.base_val) / self.base_val * 100 + + @property + def short_config(self) -> str: + return _short_config(self.config, self.device) + + +@dataclass +class ConfigAgg: + base_agg: float + base_count: int + head_agg: float + head_count: int + paired_ratio: float # gmean(head_val / base_val) over paired models + paired_count: int + + +def compute_deltas( + head_perf: list[PerfData], base_perf: list[PerfData], metric: Metric +) -> tuple[list[ModelDelta], dict[str, ConfigAgg]]: + """Join head and base on (device, config, model_name) and compute deltas. + + Returns (per_model_deltas, per_config_aggregates). + per_config_aggregates maps qualified_config -> ConfigAgg. + """ + # Build base lookup: (device, config, model_name) -> metric value + base_lookup: dict[tuple[str, str, str], float] = {} + for perf in base_perf: + for m in perf.models: + v = getattr(m, metric.field) + if v > 0: + base_lookup[(perf.device, perf.config, m.name)] = v + + deltas = [] + for perf in head_perf: + for m in perf.models: + head_val = getattr(m, metric.field) + if head_val <= 0: + continue + key = (perf.device, perf.config, m.name) + if key not in base_lookup: + continue + base_val = base_lookup[key] + deltas.append( + ModelDelta( + name=m.name, + base_val=base_val, + head_val=head_val, + config=perf.config, + device=perf.device, + ) + ) + + # Group deltas by qualified_config for paired aggregates + deltas_by_qconfig: dict[str, list[ModelDelta]] = defaultdict(list) + for d in deltas: + qc = f"{d.device}/{d.config}" if d.device else d.config + deltas_by_qconfig[qc].append(d) + + # Per-config aggregates (keyed by qualified_config for display) + config_aggs: dict[str, ConfigAgg] = {} + base_by_qconfig: dict[str, PerfData] = {p.qualified_config: p for p in base_perf} + for perf in head_perf: + qc = perf.qualified_config + if qc not in base_by_qconfig: + continue + base_perf_data = base_by_qconfig[qc] + head_agg = perf.aggregate_metric(metric) + base_agg = base_perf_data.aggregate_metric(metric) + head_count = len([m for m in perf.models if getattr(m, metric.field) > 0]) + base_count = len( + [m for m in base_perf_data.models if getattr(m, metric.field) > 0] + ) + + paired = deltas_by_qconfig.get(qc, []) + ratios = [d.head_val / d.base_val for d in paired if d.base_val > 0] + config_aggs[qc] = ConfigAgg( + base_agg=base_agg, + base_count=base_count, + head_agg=head_agg, + head_count=head_count, + paired_ratio=gmean(ratios) if ratios else 0.0, + paired_count=len(ratios), + ) + + return deltas, config_aggs + + +def print_comparison_table(config_aggs: dict[str, ConfigAgg], metric: Metric): + u = metric.unit + print(f"\n{'Config':<55} {'base':>16} {'new':>16} {'head/base':>18}") + print("─" * 108) + for config in sorted(config_aggs): + agg = config_aggs[config] + flag = "" + if agg.paired_ratio > 0: + delta_pct = (agg.paired_ratio - 1.0) * 100 + if abs(delta_pct) > RELATIVE_THRESHOLD * 100: + if metric.higher_is_better: + flag = " !!" if delta_pct < 0 else " ++" + else: + flag = " !!" if delta_pct > 0 else " ++" + print( + f" {config:<53} " + f"{agg.base_agg:>5.2f}{u} (n={agg.base_count}) " + f"{agg.head_agg:>5.2f}{u} (n={agg.head_count}) " + f"{agg.paired_ratio:>5.3f}x (n={agg.paired_count}){flag}" + ) + + +def print_regressions(deltas: list[ModelDelta], metric: Metric, top_n: int): + # For higher_is_better metrics, regression = negative delta + # For lower_is_better metrics, regression = positive delta + if metric.higher_is_better: + bad = [d for d in deltas if d.delta_pct < -RELATIVE_THRESHOLD * 100] + bad.sort(key=lambda d: d.delta_pct) + else: + bad = [d for d in deltas if d.delta_pct > RELATIVE_THRESHOLD * 100] + bad.sort(key=lambda d: d.delta_pct, reverse=True) + + if not bad: + print(f"\n No regressions (>{RELATIVE_THRESHOLD * 100:.0f}% change).") + return + + print(f"\n Regressions ({len(bad)} models, showing top {min(top_n, len(bad))}):") + for i, d in enumerate(bad[:top_n], 1): + print( + f" {i}. {d.name:<30} " + f"{d.base_val:.2f}{metric.unit} → {d.head_val:.2f}{metric.unit} " + f"{d.delta_pct:>+6.1f}% {d.short_config}" + ) + + +def render_delta_scurve( + deltas: list[ModelDelta], + metric: Metric, + term_width: int | None = None, + term_height: int | None = None, +): + if not deltas: + return + + if term_width is None or term_height is None: + sz = shutil.get_terminal_size((100, 50)) + term_width = term_width or sz.columns + term_height = term_height or sz.lines + + sorted_deltas = sorted(deltas, key=lambda d: d.delta_pct) + n = len(sorted_deltas) + max_rows = max(term_height - 5, 15) + display = subsample(sorted_deltas, max_rows) + skipped = n - len(display) + + header = f"Delta S-curve ({n} models)" + if skipped > 0: + header += f" [showing {len(display)}/{n}]" + print(f"\n {header}") + print(f" {'─' * min(len(header), term_width - 4)}") + + max_name = min(max(len(d.name) for d in display), 28) + # " name +12.3% {bar}" + prefix_len = 2 + max_name + 2 + 7 + 2 + plot_width = max(term_width - prefix_len - 1, 20) + + # Range: cap at p5/p95 to avoid outlier squishing + p5_idx = max(0, int(n * 0.05)) + p95_idx = min(n - 1, int(n * 0.95)) + range_lo = min(sorted_deltas[p5_idx].delta_pct, -10) + range_hi = max(sorted_deltas[p95_idx].delta_pct, 10) + # Ensure symmetric-ish around 0 + abs_max = max(abs(range_lo), abs(range_hi)) + range_lo = -abs_max + range_hi = abs_max + span = range_hi - range_lo + if span == 0: + span = 1 + + def pct_to_col(pct: float) -> int: + return max( + 0, min(plot_width - 1, int((pct - range_lo) / span * (plot_width - 1))) + ) + + zero_col = pct_to_col(0) + + for d in display: + name = d.name[:max_name].ljust(max_name) + col = pct_to_col(d.delta_pct) + bar = [" "] * plot_width + bar[zero_col] = "|" + if col <= zero_col: + for i in range(col, zero_col): + bar[i] = "·" + bar[zero_col] = "|" + else: + bar[zero_col] = "|" + for i in range(zero_col + 1, col + 1): + bar[i] = "·" + print(f" {name} {d.delta_pct:>+6.1f}% {''.join(bar)}") + + pad = " " * prefix_len + print(f"{pad}{' ' * zero_col}0%") + + +def print_summary_table(all_perf: list[PerfData], metric: Metric): + agg_label = metric.aggregate + print(f"\n{'Config':<65} {'Models':>6} {agg_label:>10}") + print("─" * 85) + for perf in all_perf: + n = len([m for m in perf.models if getattr(m, metric.field) > 0]) + agg = perf.aggregate_metric(metric) + print(f" {perf.qualified_config:<63} {n:>6} {agg:>8.2f}{metric.unit}") + + +def _resolve_head_runs(args) -> dict[str, int]: + """Parse run_id arg into {device: run_id} mapping. + + run_id can be: + - A single numeric run ID (device inferred from workflow) + - A branch name (resolved across all --device values, or auto-discovered) + Pre-resolved IDs can be passed via args._run_ids. + """ + # Pre-resolved (from --wait-and-summarize) + if hasattr(args, "_run_ids") and args._run_ids: + return args._run_ids + + raw = args.run_id + try: + run_id = int(raw) + # Single run ID — infer device from the workflow + meta = fetch_run_meta(run_id) + device = device_for_workflow(meta.workflow_name) + return {device: run_id} + except ValueError: + # Branch name + if args.device: + return resolve_runs(raw, args.device) + return discover_runs(raw) + + +def cmd_summary(args): + attempt = args.attempt + metric = METRICS[args.metric] + + head_run_ids = _resolve_head_runs(args) + + # Fetch head runs + auto_discovered = not hasattr(args, "_run_ids") and not args.device + head_metas: list[RunMeta] = [] + head_perf: list[PerfData] = [] + for device, run_id in list(head_run_ids.items()): + print(f"Fetching head run {run_id} ({device})...") + perf = fetch_run_perf( + run_id, + attempt, + args.no_cache, + device=device, + allow_empty=auto_discovered, + ) + if not perf: + del head_run_ids[device] + continue + head_metas.append(fetch_run_meta(run_id)) + head_perf.extend(perf) + + head_perf = filter_perf(head_perf, args) + if not head_perf: + print("No configs matched filters.") + sys.exit(1) + + head_configs = [p.qualified_config for p in head_perf] + + # Baseline comparison mode + if args.baseline and args.baseline.lower() != "none": + base_metas: list[RunMeta] = [] + base_perf: list[PerfData] = [] + + for device, head_run_id in head_run_ids.items(): + if args.baseline == "latest": + baseline_id = resolve_baseline(head_run_id) + else: + try: + baseline_id = int(args.baseline) + except ValueError: + baseline_id = resolve_run(args.baseline, device) + + print(f"Fetching baseline run {baseline_id} ({device})...") + base_data = fetch_run_perf( + baseline_id, + attempt, + args.no_cache, + device=device, + allow_empty=auto_discovered, + ) + if not base_data: + continue + base_metas.append(fetch_run_meta(baseline_id)) + base_perf.extend(base_data) + + base_perf = filter_perf(base_perf, args) + if not base_perf: + print("No baseline configs matched filters.") + sys.exit(1) + + print_run_header("HEAD", head_metas, head_configs) + print_run_header("BASE", base_metas, [p.qualified_config for p in base_perf]) + print() + + deltas, config_aggs = compute_deltas(head_perf, base_perf, metric) + if not deltas: + print("No matching models between head and baseline.") + sys.exit(1) + + print_comparison_table(config_aggs, metric) + print_regressions(deltas, metric, args.top) + render_delta_scurve(deltas, metric) + return + + # Absolute mode (no baseline) + print_run_header("RUN", head_metas, head_configs) + print() + + print_summary_table(head_perf, metric) + grouped = group_perf(head_perf, args.group_by) + for perf in grouped: + print_worst_offenders(perf, metric, args.top) + render_scurve(perf, metric, args.top) + + +CONFIG_RE = re.compile( + r"(?Pinductor_[a-z_]+?)_" + r"(?Phuggingface|timm_models|torchbench)_" + r"(?P\w+)_" + r"(?Ptraining|inference)_" + r"(?P\w+)" +) + + +def config_to_command( + config: str, + suite: str, + model: str | None = None, +) -> str | None: + """Turn a config name into a runnable benchmark command.""" + m = CONFIG_RE.match(config) + if not m: + return None + + backend_variant = m.group("backend") + dtype = m.group("dtype") + mode = m.group("mode") + runtime = m.group("device") + + # Runtime → --device flag (strip platform suffix like _x86_zen) + device_flag = runtime.split("_")[0] # "cpu_x86_zen" → "cpu" + + cmd_parts = [ + "python", + f"benchmarks/dynamo/{suite}.py", + f"--{mode}", + f"--{dtype}", + "--backend", + "inductor", + "--device", + device_flag, + ] + + if "no_cudagraphs" in backend_variant: + cmd_parts.append("--disable-cudagraphs") + if "dynamic" in backend_variant: + cmd_parts.extend(["--dynamic-shapes", "--dynamic-batch-only"]) + if "cpp_wrapper" in backend_variant: + cmd_parts.insert(0, "TORCHINDUCTOR_CPP_WRAPPER=1") + cmd_parts.append("--disable-cudagraphs") + if "freezing" in backend_variant: + cmd_parts.append("--freezing") + if "max_autotune" in backend_variant: + cmd_parts.insert(0, "TORCHINDUCTOR_MAX_AUTOTUNE=1") + if "aot_inductor" in backend_variant: + cmd_parts.append("--export-aot-inductor") + cmd_parts.append("--disable-cudagraphs") + + cmd_parts.extend(["--performance", "--cold-start-latency"]) + + if model: + cmd_parts.extend(["--only", model]) + + cmd_parts.extend(["--output", f"{config}_performance.csv"]) + + return " ".join(cmd_parts) + + +def cmd_repro(args): + run_ids = _resolve_head_runs(args) + + # Fetch perf data to discover configs + auto_discovered = not hasattr(args, "_run_ids") and not args.device + all_perf: list[PerfData] = [] + metas: list[RunMeta] = [] + for device, run_id in list(run_ids.items()): + print(f"Fetching run {run_id} ({device})...") + metas.append(fetch_run_meta(run_id)) + perf = fetch_run_perf( + run_id, + args.attempt, + no_cache=False, + device=device, + allow_empty=auto_discovered, + ) + all_perf.extend(perf) + + all_perf = filter_perf(all_perf, args) + if not all_perf: + print("No configs matched filters.") + sys.exit(1) + + print_run_header("REPRO", metas) + + configs_seen: dict[str, str] = {} # config_name → suite + for perf in all_perf: + configs_seen[perf.config] = perf.suite + + count = 0 + commands: list[str] = [] + for config in sorted(configs_seen): + suite = configs_seen[config] + cmd = config_to_command(config, suite, args.model) + if not cmd: + continue + commands.append(f"# {config}\n{cmd}") + count += 1 + + print(f"\nReproducible commands ({count} configs):\n") + for cmd in commands: + print(f"{cmd}\n") + + +PIN_DIR = Path(".ci/docker/ci_commit_pins") + + +def read_pin(name: str) -> str: + """Read a pinned commit or requirements file.""" + path = PIN_DIR / name + if not path.exists(): + return f"<{name} not found>" + return path.read_text().strip() + + +def cmd_prepare_repro(args): + suites = set() + if args.suite: + suite = SUITE_ALIASES.get(args.suite, args.suite) + suites = {suite} + else: + suites = {"huggingface", "timm_models", "torchbench"} + + torchbench_pin = read_pin("torchbench.txt") + timm_pin = read_pin("timm.txt") + hf_reqs = read_pin("huggingface-requirements.txt") + + print("# Setup commands for inductor perf benchmark suites") + print("# These mirror what CI does in the inductor-benchmarks Docker image.") + print("#") + print("# Pinned versions (commits, package versions) are read live from") + print("# .ci/docker/ci_commit_pins/") + print("# Install steps are based on:") + print("# .ci/docker/common/install_inductor_benchmark_deps.sh (build-time)") + print("# .ci/pytorch/test.sh (runtime)") + print("# If the setup process changes, check those files.") + print() + + if "huggingface" in suites: + print("# ── HuggingFace ──") + for line in hf_reqs.splitlines(): + line = line.strip() + if line and not line.startswith("#"): + print(f"pip install {line}") + print() + + if "timm_models" in suites: + print("# ── Timm ──") + print( + f"pip install git+https://github.com/huggingface/pytorch-image-models@{timm_pin}" + ) + print() + + if "torchbench" in suites: + print("# ── TorchBench ──") + print("git clone https://github.com/pytorch/benchmark torchbench") + print(f"cd torchbench && git checkout {torchbench_pin}") + print("python install.py --continue_on_fail") + print("cd ..") + print() + print("# Set PYTHONPATH so benchmark scripts find torchbench") + print("export PYTHONPATH=$(pwd)/torchbench") + print() + + print("# ── Runtime dependencies ──") + print("pip install torchvision torchaudio") + if "torchbench" in suites: + print("pip install opencv-python==4.8.0.74") + print() + + print("# ── Environment variables ──") + print("export TORCHINDUCTOR_FX_GRAPH_CACHE=True") + print("export TORCHINDUCTOR_AUTOGRAD_CACHE=True") + print() + + if not args.no_repro: + # Also print repro commands if we have a run_id + if args.run_id: + # Reuse the repro logic + repro_args = argparse.Namespace( + run_id=args.run_id, + device=args.device, + model=args.model, + suite=args.suite, + mode=args.mode, + backend=None, + dtype=None, + runtime=None, + attempt=args.attempt, + ) + print("# ── Benchmark commands ──") + cmd_repro(repro_args) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main(): + parser = argparse.ArgumentParser( + description="CLI tool for inductor perf regression runs", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Launch a perf run on your current branch (A100) + python benchmarks/dynamo/perf_cli.py launch + + # Launch on both A100 and H100 + python benchmarks/dynamo/perf_cli.py launch --device a100 h100 + + # Launch on H100 and wait for completion + python benchmarks/dynamo/perf_cli.py launch --device h100 --wait + + # Launch with only inference, no training + python benchmarks/dynamo/perf_cli.py launch --no-training --inference + + # Launch with dynamic shapes enabled + python benchmarks/dynamo/perf_cli.py launch --dynamic + + # Launch on ROCm MI300 + python benchmarks/dynamo/perf_cli.py launch --device rocm-mi300 + + # Check your branch against latest main nightly + python benchmarks/dynamo/perf_cli.py summary my-feature-branch + + # Same but on H100 + python benchmarks/dynamo/perf_cli.py summary my-feature-branch --device h100 + + # Compare across A100 and H100 in one summary + python benchmarks/dynamo/perf_cli.py summary my-feature-branch --device a100 h100 + + # Use a specific run ID instead of branch name + python benchmarks/dynamo/perf_cli.py summary 22842783236 + + # Compare against a specific baseline (run ID or branch) + python benchmarks/dynamo/perf_cli.py summary 22842783236 --baseline 22816292132 + + # Absolute metrics only (no comparison) + python benchmarks/dynamo/perf_cli.py summary 22842783236 --baseline none + + # Filter to cudagraphs training + python benchmarks/dynamo/perf_cli.py summary main --config cudagraphs --mode training + + # Show commands to reproduce a run locally for a single model + python benchmarks/dynamo/perf_cli.py repro 22842783236 --model BERT_pytorch --suite tb +""", + ) + sub = parser.add_subparsers(dest="command", required=True) + + # -- launch -- + p_launch = sub.add_parser( + "launch", help="Launch a perf regression run on your branch" + ) + p_launch.add_argument( + "--device", + nargs="+", + default=["a100"], + choices=DEVICE_CHOICES, + metavar="DEVICE", + help=f"Devices to launch (default: a100). Choices: {', '.join(DEVICE_CHOICES)}", + ) + p_launch.add_argument( + "--ref", + type=str, + default=None, + help="Git ref to benchmark (default: current branch)", + ) + p_launch.add_argument( + "--wait", action="store_true", help="Wait for all runs to complete" + ) + p_launch.add_argument( + "--wait-and-summarize", + dest="wait_and_summarize", + action="store_true", + help="Wait for all runs, then print summary vs. latest main nightly", + ) + # Workflow dispatch inputs + launch_opts = p_launch.add_argument_group("workflow options") + launch_opts.add_argument( + "--training", + action="store_true", + default=None, + help="Enable training benchmarks", + ) + launch_opts.add_argument( + "--no-training", + dest="training", + action="store_false", + help="Disable training benchmarks", + ) + launch_opts.add_argument( + "--inference", + action="store_true", + default=None, + help="Enable inference benchmarks", + ) + launch_opts.add_argument( + "--no-inference", + dest="inference", + action="store_false", + help="Disable inference benchmarks", + ) + launch_opts.add_argument( + "--cudagraphs", action="store_true", default=None, help="Enable cudagraphs" + ) + launch_opts.add_argument( + "--no-cudagraphs", + dest="cudagraphs", + action="store_false", + help="Disable cudagraphs", + ) + launch_opts.add_argument( + "--dynamic", action="store_true", default=None, help="Enable dynamic shapes" + ) + launch_opts.add_argument( + "--no-dynamic", + dest="dynamic", + action="store_false", + help="Disable dynamic shapes", + ) + launch_opts.add_argument( + "--cppwrapper", action="store_true", default=None, help="Enable cpp wrapper" + ) + launch_opts.add_argument( + "--freezing-cudagraphs", + dest="freezing_cudagraphs", + action="store_true", + default=None, + ) + launch_opts.add_argument("--aotinductor", action="store_true", default=None) + launch_opts.add_argument("--maxautotune", action="store_true", default=None) + launch_opts.add_argument( + "--default", dest="default", action="store_true", default=None + ) + launch_opts.add_argument( + "--benchmark-configs", + dest="benchmark_configs", + type=str, + default=None, + help="Override benchmark_configs input", + ) + + # -- summary -- + p_summary = sub.add_parser("summary", help="Summarize results of a perf run") + p_summary.add_argument( + "run_id", type=str, help="GitHub Actions run ID or branch name" + ) + p_summary.add_argument( + "--device", + nargs="+", + default=None, + choices=DEVICE_CHOICES, + metavar="DEVICE", + help=f"Device(s) to summarize (default: auto-discover from branch). Choices: {', '.join(DEVICE_CHOICES)}", + ) + p_summary.add_argument( + "--baseline", + type=str, + default="latest", + help="Baseline run ID, 'latest' for most recent main nightly (default), or 'none' to disable", + ) + p_summary.add_argument( + "--metric", + type=str, + default="speedup", + choices=METRIC_CHOICES, + help="Metric to display (default: speedup)", + ) + p_summary.add_argument( + "--top", + type=int, + default=5, + help="Number of worst offenders to show (default: 5)", + ) + # Filters + filters = p_summary.add_argument_group("filters") + filters.add_argument( + "--config", + type=str, + default=None, + help="Regex to filter config names (e.g. 'cudagraphs', 'dynamic')", + ) + filters.add_argument( + "--suite", type=str, default=None, help="Filter to suite: hf, timm, tb" + ) + filters.add_argument( + "--mode", + type=str, + choices=["training", "inference"], + default=None, + help="Filter to training or inference", + ) + filters.add_argument( + "--backend", + type=str, + default=None, + help="Regex to filter backend (e.g. 'cudagraphs', 'dynamic')", + ) + filters.add_argument( + "--dtype", + type=str, + default=None, + help="Filter to dtype (e.g. amp, float16, bfloat16)", + ) + filters.add_argument( + "--runtime", + type=str, + default=None, + help="Filter to runtime (e.g. cuda, cpu, xpu)", + ) + p_summary.add_argument( + "--group-by", + dest="group_by", + nargs="+", + default=None, + choices=GROUP_CHOICES, + metavar="KEY", + help=f"Group S-curves by key(s) (default: single combined). Choices: {', '.join(GROUP_CHOICES)}", + ) + p_summary.add_argument( + "--attempt", type=int, default=1, help="Run attempt number (default: 1)" + ) + p_summary.add_argument( + "--no-cache", + dest="no_cache", + action="store_true", + default=False, + help="Re-download artifacts even if cached", + ) + + # -- repro -- + p_repro = sub.add_parser("repro", help="Reproduce a remote perf run locally") + p_repro.add_argument( + "run_id", type=str, help="GitHub Actions run ID or branch name" + ) + p_repro.add_argument( + "--device", + nargs="+", + default=None, + choices=DEVICE_CHOICES, + metavar="DEVICE", + help="Device(s) (default: auto-discover from branch)", + ) + p_repro.add_argument( + "--model", + type=str, + default=None, + help="Run only this model (e.g. BERT_pytorch)", + ) + p_repro.add_argument( + "--suite", type=str, default=None, help="Filter to suite (hf, timm, tb)" + ) + p_repro.add_argument( + "--mode", + type=str, + choices=["training", "inference"], + default=None, + help="Filter to mode", + ) + p_repro.add_argument( + "--backend", + type=str, + default=None, + help="Regex to filter backend (e.g. 'cudagraphs', 'dynamic')", + ) + p_repro.add_argument( + "--dtype", type=str, default=None, help="Filter to dtype (e.g. amp, bfloat16)" + ) + p_repro.add_argument( + "--runtime", type=str, default=None, help="Filter to runtime (e.g. cuda, cpu)" + ) + p_repro.add_argument( + "--attempt", type=int, default=1, help="Run attempt number (default: 1)" + ) + + # -- prepare-repro -- + p_prep = sub.add_parser( + "prepare-repro", help="Show setup commands to prepare benchmark suites locally" + ) + p_prep.add_argument( + "run_id", + nargs="?", + type=str, + default=None, + help="Optional: run ID or branch to also show benchmark commands", + ) + p_prep.add_argument( + "--device", + nargs="+", + default=None, + choices=DEVICE_CHOICES, + metavar="DEVICE", + help="Device(s) for benchmark commands", + ) + p_prep.add_argument( + "--suite", + type=str, + default=None, + help="Only show setup for this suite (hf, timm, tb)", + ) + p_prep.add_argument( + "--mode", + type=str, + choices=["training", "inference"], + default=None, + help="Filter benchmark commands to mode", + ) + p_prep.add_argument( + "--model", type=str, default=None, help="Filter benchmark commands to model" + ) + p_prep.add_argument( + "--no-repro", + dest="no_repro", + action="store_true", + help="Only show setup, skip benchmark commands", + ) + p_prep.add_argument( + "--attempt", type=int, default=1, help="Run attempt number (default: 1)" + ) + + args = parser.parse_args() + + if args.command == "launch": + cmd_launch(args) + elif args.command == "summary": + cmd_summary(args) + elif args.command == "repro": + cmd_repro(args) + elif args.command == "prepare-repro": + cmd_prepare_repro(args) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv b/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv index ccd7d1acd7dc7..375c73eb6e917 100644 --- a/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv +++ b/benchmarks/dynamo/pr_time_benchmarks/expected_results.csv @@ -1,8 +1,8 @@ -add_loop_eager,compile_time_instruction_count,3347000000,0.1 +add_loop_eager,compile_time_instruction_count,3479000000,0.1 -add_loop_eager_dynamic,compile_time_instruction_count,4578000000,0.1 +add_loop_eager_dynamic,compile_time_instruction_count,4779000000,0.1 @@ -10,7 +10,7 @@ add_loop_inductor,compile_time_instruction_count,29660000000,0.1 -add_loop_inductor_dynamic_gpu,compile_time_instruction_count,37280000000,0.1 +add_loop_inductor_dynamic_gpu,compile_time_instruction_count,38150000000,0.1 @@ -18,11 +18,11 @@ add_loop_inductor_gpu,compile_time_instruction_count,26140000000,0.1 -basic_modules_ListOfLinears_eager,compile_time_instruction_count,898000000,0.1 +basic_modules_ListOfLinears_eager,compile_time_instruction_count,932400000,0.1 -basic_modules_ListOfLinears_inductor,compile_time_instruction_count,15240000000,0.1 +basic_modules_ListOfLinears_inductor,compile_time_instruction_count,15640000000,0.1 @@ -34,19 +34,19 @@ basic_modules_ListOfLinears_inductor_gpu,compile_time_instruction_count,11090000 -update_hint_regression,compile_time_instruction_count,1548000000,0.1 +update_hint_regression,compile_time_instruction_count,1609000000,0.1 -sum_floordiv_regression,compile_time_instruction_count,3477000000,0.1 +sum_floordiv_regression,compile_time_instruction_count,3641000000,0.1 -symint_sum,compile_time_instruction_count,3112000000,0.1 +symint_sum,compile_time_instruction_count,3210000000,0.1 -symint_sum_loop,compile_time_instruction_count,4305000000,0.1 +symint_sum_loop,compile_time_instruction_count,4448000000,0.1 @@ -54,15 +54,15 @@ aotdispatcher_inference_nosubclass_cpu,compile_time_instruction_count,1793000000 -aotdispatcher_inference_subclass_cpu,compile_time_instruction_count,4988000000,0.1 +aotdispatcher_inference_subclass_cpu,compile_time_instruction_count,5109000000,0.1 -aotdispatcher_partitioner_cpu,compile_time_instruction_count,7858000000,0.1 +aotdispatcher_partitioner_cpu,compile_time_instruction_count,8318000000,0.1 -aotdispatcher_partitioner_cpu2,compile_time_instruction_count,1834000000,0.1 +aotdispatcher_partitioner_cpu2,compile_time_instruction_count,1937000000,0.1 @@ -70,23 +70,23 @@ aotdispatcher_training_nosubclass_cpu,compile_time_instruction_count,3152000000, -aotdispatcher_training_subclass_cpu,compile_time_instruction_count,8021000000,0.1 +aotdispatcher_training_subclass_cpu,compile_time_instruction_count,8317000000,0.1 -mm_loop_inductor_gpu,compile_time_instruction_count,4597000000,0.1 +mm_loop_inductor_gpu,compile_time_instruction_count,4707000000,0.1 -mm_loop_inductor_dynamic_gpu,compile_time_instruction_count,9051000000,0.1 +mm_loop_inductor_dynamic_gpu,compile_time_instruction_count,9270000000,0.1 -basic_NestedModule_eager,compile_time_instruction_count,5728000000,0.1 +basic_NestedModule_eager,compile_time_instruction_count,5866000000,0.1 -basic_InlineMod_eager,compile_time_instruction_count,7701000000,0.1 +basic_InlineMod_eager,compile_time_instruction_count,7863000000,0.1 @@ -102,11 +102,11 @@ dtensor_dispatch_collectives,instruction_count,28890000,0.1 -dtensor_dispatch_add_backward,instruction_count,346201,0.1 +dtensor_dispatch_add_backward,instruction_count,360100,0.1 -dtensor_dispatch_inplace,instruction_count,56530,0.1 +dtensor_dispatch_inplace,instruction_count,61490,0.1 @@ -118,4 +118,4 @@ dtensor_dispatch_random,instruction_count,592800,0.1 -dtensor_dispatch_custom_handler,instruction_count,35020,0.1 +dtensor_dispatch_custom_handler,instruction_count,36050,0.1 diff --git a/benchmarks/dynamo/timm_models.py b/benchmarks/dynamo/timm_models.py index 4855e1cf1e13d..09e7647537a4a 100755 --- a/benchmarks/dynamo/timm_models.py +++ b/benchmarks/dynamo/timm_models.py @@ -56,54 +56,6 @@ def pip_install(package): TIMM_MODELS[model_name] = int(batch_size) -# TODO - Figure out the reason of cold start memory spike - -BATCH_SIZE_DIVISORS = { - "beit_base_patch16_224": 2, - "deit_base_distilled_patch16_224": 2, - "gluon_xception65": 2, - "mobilevit_s": 2, - "swin_base_patch4_window7_224": 2, -} - -REQUIRE_HIGHER_TOLERANCE = { - "inception_v3", - "mobilenetv3_large_100", -} - -REQUIRE_HIGHER_TOLERANCE_FP16_XPU = { - "botnet26t_256", -} - -REQUIRE_HIGHER_TOLERANCE_AMP = {} - -REQUIRE_EVEN_HIGHER_TOLERANCE = { - "deit_base_distilled_patch16_224", - "vit_base_patch16_siglip_256", -} - -# These models need higher tolerance in MaxAutotune mode -REQUIRE_EVEN_HIGHER_TOLERANCE_MAX_AUTOTUNE = {} - -REQUIRE_HIGHER_TOLERANCE_FOR_FREEZING = { - "adv_inception_v3", -} - -SCALED_COMPUTE_LOSS = { - "mobilevit_s", -} - -FORCE_AMP_FOR_FP16_BF16_MODELS = {} - -SKIP_ACCURACY_CHECK_AS_EAGER_NON_DETERMINISTIC_MODELS = {} - -REQUIRE_LARGER_MULTIPLIER_FOR_SMALLER_TENSOR = { - "inception_v3", - "mobilenetv3_large_100", - "vit_base_patch14_dinov2.lvd142m", -} - - def refresh_model_names(): import glob @@ -200,6 +152,22 @@ def _config(self): def _skip(self): return self._config["skip"] + @property + def _batch_size(self): + return self._config["batch_size"] + + @property + def _tolerance(self): + return self._config["tolerance"] + + @property + def _accuracy(self): + return self._config["accuracy"] + + @property + def _require_larger_multiplier_for_smaller_tensor(self): + return self._config["require_larger_multiplier_for_smaller_tensor"] + @property def skip_models_for_cpu(self): return self._skip["device"]["cpu"] @@ -214,7 +182,7 @@ def skip_models(self): @property def force_amp_for_fp16_bf16_models(self): - return FORCE_AMP_FOR_FP16_BF16_MODELS + return self._config["dtype"]["force_amp_for_fp16_bf16_models"] @property def force_fp16_for_bf16_models(self): @@ -227,7 +195,7 @@ def get_output_amp_train_process_func(self): @property def skip_accuracy_check_as_eager_non_deterministic(self): if self.args.accuracy and self.args.training: - return SKIP_ACCURACY_CHECK_AS_EAGER_NON_DETERMINISTIC_MODELS + return self._accuracy["skip"]["eager_not_deterministic"] return set() @property @@ -281,9 +249,10 @@ def load_model( input_size = data_config["input_size"] recorded_batch_size = TIMM_MODELS[model_name] - if model_name in BATCH_SIZE_DIVISORS: + batch_size_divisors = self._batch_size["divisors"] + if model_name in batch_size_divisors: recorded_batch_size = max( - int(recorded_batch_size / BATCH_SIZE_DIVISORS[model_name]), 1 + int(recorded_batch_size / batch_size_divisors[model_name]), 1 ) batch_size = batch_size or recorded_batch_size @@ -305,7 +274,7 @@ def load_model( self.loss = torch.nn.CrossEntropyLoss().to(device) - if model_name in SCALED_COMPUTE_LOSS: + if model_name in self._config["scaled_compute_loss"]: self.compute_loss = self.scaled_compute_loss if is_training and not use_eval_mode: @@ -313,7 +282,7 @@ def load_model( else: model.eval() - self.validate_model(model, example_inputs) + self.validate_model(model_name, model, example_inputs) return device, model_name, model, example_inputs, batch_size @@ -341,34 +310,34 @@ def pick_grad(self, name, is_training): return torch.no_grad() def use_larger_multiplier_for_smaller_tensor(self, name): - return name in REQUIRE_LARGER_MULTIPLIER_FOR_SMALLER_TENSOR + return name in self._require_larger_multiplier_for_smaller_tensor def get_tolerance_and_cosine_flag(self, is_training, current_device, name): cosine = self.args.cosine tolerance = 1e-3 - if self.args.freezing and name in REQUIRE_HIGHER_TOLERANCE_FOR_FREEZING: + if self.args.freezing and name in self._tolerance["freezing"]: # the conv-batchnorm fusion used under freezing may cause relatively - # large numerical difference. We need are larger tolerance. + # large numerical difference. We need a larger tolerance. # Check https://github.com/pytorch/pytorch/issues/120545 for context tolerance = 8 * 1e-2 if is_training: from torch._inductor import config as inductor_config - if name == "beit_base_patch16_224": + if name in self._tolerance["highest_training"]: tolerance = 16 * 1e-2 - elif name in REQUIRE_EVEN_HIGHER_TOLERANCE or ( + elif name in self._tolerance["even_higher"] or ( inductor_config.max_autotune - and name in REQUIRE_EVEN_HIGHER_TOLERANCE_MAX_AUTOTUNE + and name in self._tolerance["even_higher_max_autotune"] ): tolerance = 8 * 1e-2 - elif name in REQUIRE_HIGHER_TOLERANCE or ( - self.args.amp and name in REQUIRE_HIGHER_TOLERANCE_AMP + elif name in self._tolerance["higher_training"] or ( + self.args.amp and name in self._tolerance["higher_amp"] ): tolerance = 4 * 1e-2 elif ( - name in REQUIRE_HIGHER_TOLERANCE_FP16_XPU + name in self._tolerance["higher_fp16_xpu"] and self.args.float16 and current_device == "xpu" ): diff --git a/benchmarks/dynamo/timm_models.yaml b/benchmarks/dynamo/timm_models.yaml index 1650a87500537..3e94903791f93 100644 --- a/benchmarks/dynamo/timm_models.yaml +++ b/benchmarks/dynamo/timm_models.yaml @@ -11,3 +11,62 @@ skip: - dm_nfnet_f0 - nfnet_l0 - visformer_small + + +# TODO - Figure out the reason of cold start memory spike +batch_size: + divisors: + beit_base_patch16_224: 2 + deit_base_distilled_patch16_224: 2 + gluon_xception65: 2 + mobilevit_s: 2 + swin_base_patch4_window7_224: 2 + + +tolerance: + higher_training: + - inception_v3 + - mobilenetv3_large_100 + + higher_fp16_xpu: + - botnet26t_256 + + higher_amp: [] + + even_higher: + - deit_base_distilled_patch16_224 + - vit_base_patch16_siglip_256 + + # These models need higher tolerance in MaxAutotune mode + even_higher_max_autotune: [] + + # beit_base_patch16_224 needs 16 * 1e-2 tolerance in training + highest_training: + - beit_base_patch16_224 + + freezing: + # the conv-batchnorm fusion used under freezing may cause relatively + # large numerical difference. We need a larger tolerance. + # Check https://github.com/pytorch/pytorch/issues/120545 for context + - adv_inception_v3 + + cosine: [] + + +scaled_compute_loss: + - mobilevit_s + + +require_larger_multiplier_for_smaller_tensor: + - inception_v3 + - mobilenetv3_large_100 + - vit_base_patch14_dinov2.lvd142m + + +dtype: + force_amp_for_fp16_bf16_models: [] + + +accuracy: + skip: + eager_not_deterministic: [] diff --git a/benchmarks/dynamo/torchbench.py b/benchmarks/dynamo/torchbench.py index 8b79aef6596ff..cbe575f926d22 100755 --- a/benchmarks/dynamo/torchbench.py +++ b/benchmarks/dynamo/torchbench.py @@ -381,7 +381,7 @@ def load_model( ): model.config.use_cache = False - self.validate_model(model, example_inputs) + self.validate_model(benchmark.name, model, example_inputs) return device, benchmark.name, model, example_inputs, batch_size def iter_model_names(self, args): diff --git a/benchmarks/operator_benchmark/benchmark_runner.py b/benchmarks/operator_benchmark/benchmark_runner.py index 6568cf9bf3ee6..4c96b40b8db47 100644 --- a/benchmarks/operator_benchmark/benchmark_runner.py +++ b/benchmarks/operator_benchmark/benchmark_runner.py @@ -16,6 +16,7 @@ description="Run microbenchmarks.", formatter_class=argparse.ArgumentDefaultsHelpFormatter, conflict_handler="resolve", + allow_abbrev=False, ) diff --git a/benchmarks/serialization/export_save_linear_bench.py b/benchmarks/serialization/export_save_linear_bench.py new file mode 100644 index 0000000000000..6f3f9360e9896 --- /dev/null +++ b/benchmarks/serialization/export_save_linear_bench.py @@ -0,0 +1,287 @@ +import argparse +import gc +import statistics +import tempfile +import time +from dataclasses import dataclass +from pathlib import Path + +import torch + + +DEFAULT_NUM_PARAMS = ( + 1_000_000, + 10_000_000, + 100_000_000, + 1_000_000_000, + 2_000_000_000, + 3_000_000_000, + 4_000_000_000, +) +NUM_LAYERS = 5 +DEFAULT_REPEATS = 5 + + +class LinearModel(torch.nn.Module): + def __init__(self, hidden_size: int, *, dtype: torch.dtype) -> None: + super().__init__() + self.layers = torch.nn.ModuleList( + [ + torch.nn.Linear( + hidden_size, + hidden_size, + bias=True, + dtype=dtype, + ) + for _ in range(NUM_LAYERS) + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for layer in self.layers: + x = layer(x) + return x + + +@dataclass +class BenchmarkResult: + num_params: int + median_save_ms: float | None + status: str + error_detail: str | None = None + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Benchmark torch.export.save for a five-layer LinearModel across " + "a range of parameter counts." + ) + ) + parser.add_argument( + "--num-params", + type=int, + nargs="+", + default=list(DEFAULT_NUM_PARAMS), + help=( + "Target total parameter counts to benchmark. The script derives the " + "nearest hidden size for each target." + ), + ) + parser.add_argument( + "--repeats", + type=int, + default=DEFAULT_REPEATS, + help="Number of torch.export.save timings to collect per parameter-count case.", + ) + parser.add_argument( + "--batch-size", + type=int, + default=0, + help=( + "Batch size for the example input passed to torch.export.export. " + "The default of 0 keeps the benchmark focused on serialization cost." + ), + ) + parser.add_argument( + "--dtype", + choices=("float32", "float16", "bfloat16"), + default="float32", + help="Parameter and example-input dtype.", + ) + return parser.parse_args() + + +def _dtype_from_name(dtype_name: str) -> torch.dtype: + return { + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + }[dtype_name] + + +def _parameter_count(hidden_size: int) -> int: + return NUM_LAYERS * (hidden_size * hidden_size + hidden_size) + + +def _hidden_size_from_num_params(num_params: int) -> int: + if num_params <= 0: + raise ValueError("--num-params values must be positive") + + root = (-1.0 + (1.0 + 4.0 * num_params / NUM_LAYERS) ** 0.5) / 2.0 + base_hidden_size = max(1, int(root)) + candidates = [ + candidate + for candidate in ( + base_hidden_size - 1, + base_hidden_size, + base_hidden_size + 1, + base_hidden_size + 2, + ) + if candidate >= 1 + ] + return min( + candidates, + key=lambda hidden_size: abs(_parameter_count(hidden_size) - num_params), + ) + + +def _format_num_params(num_params: int) -> str: + for threshold, suffix in ( + (1_000_000_000, "B"), + (1_000_000, "M"), + (1_000, "K"), + ): + if num_params >= threshold: + value = num_params / threshold + return f"{value:.1f}".rstrip("0").rstrip(".") + suffix + return str(num_params) + + +def _measure_save_times_ms( + exported_program: torch.export.ExportedProgram, + repeats: int, +) -> list[float]: + times_ms: list[float] = [] + with tempfile.TemporaryDirectory(prefix="export_save_bench_") as temp_dir: + temp_path = Path(temp_dir) + for iteration in range(repeats): + save_path = temp_path / f"exported_program_{iteration}.pt2" + start_time = time.perf_counter() + torch.export.save(exported_program, save_path) + elapsed_ms = (time.perf_counter() - start_time) * 1000.0 + times_ms.append(elapsed_ms) + save_path.unlink(missing_ok=True) + return times_ms + + +def _status_for_exception(exc: BaseException) -> tuple[str, str]: + message = ( + str(exc).strip().splitlines()[0] if str(exc).strip() else type(exc).__name__ + ) + lowered = message.lower() + if ( + isinstance(exc, MemoryError) + or "out of memory" in lowered + or "can't allocate memory" in lowered + ): + return "oom", message + return "error", message + + +def _run_case( + target_num_params: int, + *, + repeats: int, + batch_size: int, + dtype: torch.dtype, +) -> BenchmarkResult: + exported_program = None + hidden_size = _hidden_size_from_num_params(target_num_params) + num_params = _parameter_count(hidden_size) + try: + with torch.no_grad(): + model = LinearModel(hidden_size, dtype=dtype).eval() + model.requires_grad_(False) + example_input = torch.zeros(batch_size, hidden_size, dtype=dtype) + exported_program = torch.export.export( + model, + (example_input,), + strict=True, + ) + + save_times_ms = _measure_save_times_ms(exported_program, repeats) + return BenchmarkResult( + num_params=num_params, + median_save_ms=statistics.median(save_times_ms), + status="ok", + ) + except (MemoryError, RuntimeError) as exc: + status, detail = _status_for_exception(exc) + return BenchmarkResult( + num_params=num_params, + median_save_ms=None, + status=status, + error_detail=detail, + ) + finally: + del exported_program + del example_input + del model + gc.collect() + + +def _format_table(results: list[BenchmarkResult]) -> str: + headers = ("num_params", "median_save_ms", "status") + rows = [ + ( + _format_num_params(result.num_params), + ( + f"{result.median_save_ms:.3f}" + if result.median_save_ms is not None + else "n/a" + ), + result.status, + ) + for result in results + ] + widths = [ + max(len(header), *(len(row[column]) for row in rows)) + for column, header in enumerate(headers) + ] + + def _format_row(row: tuple[str, ...]) -> str: + return " | ".join( + [ + row[0].rjust(widths[0]), + row[1].rjust(widths[1]), + row[2].ljust(widths[2]), + ] + ) + + lines = [ + _format_row(headers), + "-+-".join("-" * width for width in widths), + ] + lines.extend(_format_row(row) for row in rows) + return "\n".join(lines) + + +def main() -> None: + args = _parse_args() + dtype = _dtype_from_name(args.dtype) + + print("Benchmarking torch.export.save for a 5-layer LinearModel") + print(f"dtype={args.dtype}, batch_size={args.batch_size}, repeats={args.repeats}") + + results = [ + _run_case( + num_params, + repeats=args.repeats, + batch_size=args.batch_size, + dtype=dtype, + ) + for num_params in args.num_params + ] + + print() + print( + "Note: This benchmark is highly sensitive to disk performance and OS " + "behavior. Results for smaller parameter counts can be noisy or flaky." + ) + print() + print(_format_table(results)) + + failures = [result for result in results if result.error_detail is not None] + if failures: + print() + print("failed cases:") + for result in failures: + print( + f"- num_params={_format_num_params(result.num_params)}: " + f"{result.error_detail}" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/static_runtime/test_static_runtime.cc b/benchmarks/static_runtime/test_static_runtime.cc index 69d925675d478..cf8f333bc2c25 100644 --- a/benchmarks/static_runtime/test_static_runtime.cc +++ b/benchmarks/static_runtime/test_static_runtime.cc @@ -1074,6 +1074,10 @@ TEST(StaticRuntime, NanToNum) { } TEST(StaticRuntime, Stack) { +#if defined(__aarch64__) || defined(_M_ARM64) + // See https://github.com/pytorch/pytorch/issues/178522. + GTEST_SKIP() << "Skipping StaticRuntime.Stack on AArch64."; +#endif const auto stack_dim = R"JIT( def forward(self, a: Tensor, b: Tensor, dim: int): inputs = [a] @@ -2307,6 +2311,10 @@ TEST(StaticRuntime, Append) { } TEST(StaticRuntime, QuantizedLinear) { +#if defined(__aarch64__) || defined(_M_ARM64) + // See https://github.com/pytorch/pytorch/issues/178522. + GTEST_SKIP() << "Skipping QuantizedLinear on AArch64."; +#endif const std::string quantize_script = R"IR( graph(%input: Tensor, %weights: Tensor): %scale: float = prim::Constant[value=1.]() diff --git a/buckbuild.bzl b/buckbuild.bzl index 7b1ad284d51cc..4ad6de567c4b3 100644 --- a/buckbuild.bzl +++ b/buckbuild.bzl @@ -71,15 +71,15 @@ def read_bool(section, field, default, required = True): else: fail("`{}:{}`: no value set".format(section, field)) -def _is_build_mode_dev(): - if is_production_build_android(): - # Android Prod builds - return False +def _select_if_build_mode_dev(dev_value, default = []): if is_production_build_ios() or is_profile_build_ios(): - # iOS Prod builds - return False + return default + + return select({ + "DEFAULT": default, + "ovr_config//build_mode:optimization[dev]": dev_value, + }) - return True def _get_enable_lightweight_dispatch(): return read_bool("pt", "enable_lightweight_dispatch", False) @@ -93,10 +93,24 @@ def get_enable_mobile_dispatch_keys_trimming(): def get_disable_per_op_profiling(): return read_bool("pt", "disable_per_op_profiling", True) -def get_strip_error_messages(): +def strip_error_messages_select(value, default = []): if IS_OSS: - return True # always strip in OSS CI to expose potential issues - return read_bool("pt", "strip_error_messages", not _is_build_mode_dev()) + return value # always strip in OSS CI to expose potential issues + strip_error = read_bool("pt", "strip_error_messages", default = None, required = False) + + if strip_error == None: + + if is_production_build_ios() or is_profile_build_ios(): + return value + + return select({ + "DEFAULT": default, + "ovr_config//build_mode:optimization[opt]": value, + }) + + if strip_error: + return value + return default def get_disable_warn(): return read_bool("pt", "disable_warn", False) @@ -249,9 +263,7 @@ _COMMON_PREPROCESSOR_FLAGS = [ "-DNO_EXPORT", ] + ( ["-DC10_MOBILE_TRIM_DISPATCH_KEYS"] if get_enable_mobile_dispatch_keys_trimming() else [] -) + ( - ["-DSTRIP_ERROR_MESSAGES"] if get_strip_error_messages() else [] -) + ( +) + strip_error_messages_select(["-DSTRIP_ERROR_MESSAGES"]) + ( ["-DDISABLE_WARN"] if get_disable_warn() else [] ) @@ -282,9 +294,9 @@ def get_aten_preprocessor_flags(): "-DUSE_RUY_QMATMUL", ] if get_disable_per_op_profiling(): - ATEN_PREPROCESSOR_FLAGS.append("-DPYTORCH_DISABLE_PER_OP_PROFILING") + ATEN_PREPROCESSOR_FLAGS += ["-DPYTORCH_DISABLE_PER_OP_PROFILING"] if _get_enable_record_kernel_dtype(): - ATEN_PREPROCESSOR_FLAGS.append("-DENABLE_RECORD_KERNEL_FUNCTION_DTYPE") + ATEN_PREPROCESSOR_FLAGS += ["-DENABLE_RECORD_KERNEL_FUNCTION_DTYPE"] return ATEN_PREPROCESSOR_FLAGS def get_pt_preprocessor_flags(): @@ -295,8 +307,7 @@ def get_pt_preprocessor_flags(): "-DNO_CUDNN_DESTROY_HANDLE", ] - if _is_build_mode_dev(): - PT_PREPROCESSOR_FLAGS.append("-DENABLE_PYTORCH_NON_PRODUCTION_BUILDS") + PT_PREPROCESSOR_FLAGS += _select_if_build_mode_dev(["-DENABLE_PYTORCH_NON_PRODUCTION_BUILDS"]) return PT_PREPROCESSOR_FLAGS # This needs to be kept in sync with https://github.com/pytorch/pytorch/blob/release/1.9/torchgen/gen.py#L892 @lint-ignore @@ -905,6 +916,7 @@ def define_buck_targets( ("aten/src", "ATen/ops/*.h"), # ATen Base ("aten/src", "ATen/*.h"), + ("aten/src", "ATen/accelerator/*.h"), ("aten/src", "ATen/cpu/**/*.h"), ("aten/src", "ATen/detail/*.h"), ("aten/src", "ATen/functorch/**/*.h"), diff --git a/build_variables.bzl b/build_variables.bzl index fab555b9c1714..91c5923b6cd76 100644 --- a/build_variables.bzl +++ b/build_variables.bzl @@ -691,7 +691,7 @@ libtorch_lite_cmake_sources = sorted( torch_mobile_core, ) -libtorch_cmake_sources = libtorch_core_sources + libtorch_core_jit_sources + libtorch_nativert_sources +libtorch_cmake_sources = libtorch_core_sources + libtorch_core_jit_sources libtorch_extra_sources = libtorch_core_jit_sources + [ "torch/csrc/autograd/TraceTypeManual.cpp", @@ -742,6 +742,7 @@ libtorch_cuda_core_sources = [ "torch/csrc/cuda/shim_common.cpp", "torch/csrc/inductor/aoti_runner/model_container_runner_cuda.cpp", "torch/csrc/inductor/aoti_torch/shim_cuda.cpp", + "torch/csrc/inductor/inductor_ops_gpu.cpp", "torch/csrc/jit/codegen/fuser/cuda/fused_kernel.cpp", "torch/csrc/profiler/stubs/cuda.cpp", "torch/csrc/autograd/functions/comm.cpp", @@ -776,6 +777,7 @@ libtorch_cuda_distributed_extra_sources = [ "torch/csrc/distributed/c10d/symm_mem/CudaDMAConnectivity.cpp", "torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.cu", "torch/csrc/distributed/c10d/symm_mem/nccl_extension.cu", + "torch/csrc/distributed/c10d/symm_mem/ops/nccl_reduce_scatter_offset.cu", "torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cpp", "torch/csrc/distributed/c10d/symm_mem/intra_node_comm.cu", "torch/csrc/distributed/c10d/symm_mem/cuda_mem_pool.cpp", @@ -1159,6 +1161,7 @@ aten_cpu_source_non_codegen_list = [ "aten/src/ATen/LegacyVmapMode.cpp", "aten/src/ATen/LegacyVmapTransforms.cpp", "aten/src/ATen/core/BackendSelectFallbackKernel.cpp", + "aten/src/ATen/core/CachingHostAllocator.cpp", "aten/src/ATen/core/DeprecatedTypeProperties.cpp", "aten/src/ATen/core/DeprecatedTypePropertiesRegistry.cpp", "aten/src/ATen/core/Dict.cpp", @@ -1172,7 +1175,6 @@ aten_cpu_source_non_codegen_list = [ "aten/src/ATen/core/Tensor.cpp", "aten/src/ATen/core/VariableFallbackKernel.cpp", "aten/src/ATen/core/VariableHooksInterface.cpp", - "aten/src/ATen/core/Vitals.cpp", "aten/src/ATen/core/boxing/KernelFunction.cpp", "aten/src/ATen/core/custom_class.cpp", "aten/src/ATen/core/dispatch/DispatchKeyExtractor.cpp", @@ -1503,6 +1505,7 @@ aten_native_source_non_codegen_list = [ # "aten/src/ATen/native/UpSample.cpp", "aten/src/ATen/native/UpSampleBicubic2d.cpp", "aten/src/ATen/native/UpSampleBilinear2d.cpp", + "aten/src/ATen/native/UpSampleLanczos2d.cpp", "aten/src/ATen/native/UpSampleLinear1d.cpp", "aten/src/ATen/native/UpSampleNearest1d.cpp", "aten/src/ATen/native/UpSampleNearest2d.cpp", diff --git a/c10/core/CachingDeviceAllocator.h b/c10/core/CachingDeviceAllocator.h index 6800c79ab09ff..2e26e4091abe7 100644 --- a/c10/core/CachingDeviceAllocator.h +++ b/c10/core/CachingDeviceAllocator.h @@ -56,6 +56,9 @@ struct DeviceStats { // un-mapped and free memory. int64_t num_device_free = 0; + // COUNT: total number of allocations rejected by OOM preemption policy + int64_t num_oom_rejections = 0; + // SIZE: maximum block size that is allowed to be split. int64_t max_split_size = 0; }; diff --git a/c10/core/ScalarType.h b/c10/core/ScalarType.h index 040c6abb7d8e2..dc9d168f053e7 100644 --- a/c10/core/ScalarType.h +++ b/c10/core/ScalarType.h @@ -54,6 +54,31 @@ inline size_t elementSize(ScalarType t) { #undef CASE_ELEMENTSIZE_CASE } +inline ScalarType opaqueScalarType(ScalarType t) { + auto esize = elementSize(t); + ScalarType result; + switch (esize) { + case 1: + result = kByte; + break; + case 2: + result = kUInt16; + break; + case 4: + result = kUInt32; + break; + case 8: + result = kUInt64; + break; + case 16: + result = kComplexDouble; + break; + default: + TORCH_CHECK(false, "Unknown ScalarType"); + } + return result; +} + inline bool isIntegralType(ScalarType t, bool includeBool) { bool isIntegral = (t == ScalarType::Byte || t == ScalarType::Char || t == ScalarType::Int || diff --git a/c10/core/StorageImpl.h b/c10/core/StorageImpl.h index 0cd3a8fd5db0f..8479256c7ca86 100644 --- a/c10/core/StorageImpl.h +++ b/c10/core/StorageImpl.h @@ -312,6 +312,14 @@ struct C10_API StorageImpl : public c10::intrusive_ptr_target { refresh_has_data_ptr_check(); } + void clear_data_ptr_access_error_msg_() { + throw_on_immutable_data_ptr_ = false; + if (extra_meta_) { + extra_meta_->custom_data_ptr_error_msg_ = std::nullopt; + } + refresh_has_data_ptr_check(); + } + void set_throw_on_mutable_data_ptr() { throw_on_mutable_data_ptr_ = true; refresh_has_data_ptr_check(); diff --git a/c10/core/Stream.cpp b/c10/core/Stream.cpp index 45339719a91b1..afb8f525dabe4 100644 --- a/c10/core/Stream.cpp +++ b/c10/core/Stream.cpp @@ -22,6 +22,12 @@ void Stream::synchronize() const { impl.synchronizeStream(*this); } +// Return whether this stream is currently under graph capturing mode. +bool Stream::is_capturing() const { + impl::VirtualGuardImpl impl{device_.type()}; + return impl.isStreamCapturing(*this); +} + // Not very parsable, but I don't know a good compact syntax for streams. // Feel free to change this into something more compact if needed. std::ostream& operator<<(std::ostream& stream, const Stream& s) { diff --git a/c10/core/Stream.h b/c10/core/Stream.h index 3890faa05d9c4..08eb0eed3c6e1 100644 --- a/c10/core/Stream.h +++ b/c10/core/Stream.h @@ -134,6 +134,10 @@ class C10_API Stream final { // on this stream has completed running on the device. void synchronize() const; + // Return the stream is currently recording work for graph capture. True while + // the stream is in capture mode, false otherwise. + bool is_capturing() const; + // The purpose of this function is to more conveniently permit binding // of Stream to and from Python. Without packing, I have to setup a whole // class with two fields (device and stream id); with packing I can just diff --git a/c10/core/SymbolicShapeMeta.h b/c10/core/SymbolicShapeMeta.h index 0820038968a8e..59c3147c878f7 100644 --- a/c10/core/SymbolicShapeMeta.h +++ b/c10/core/SymbolicShapeMeta.h @@ -141,7 +141,7 @@ class C10_API SymbolicShapeMeta { available_.fetch_or(is_contiguous_avail); } void assume_channels_last_contiguous(SymBool val = true) { - is_contiguous_ = std::move(val); + is_channels_last_contiguous_ = std::move(val); available_.fetch_or(is_channels_last_contiguous_avail); } void assume_channels_last_3d_contiguous(SymBool val = true) { diff --git a/c10/core/impl/DeviceGuardImplInterface.cpp b/c10/core/impl/DeviceGuardImplInterface.cpp index 428ea63c04151..6d69c31e3e1e6 100644 --- a/c10/core/impl/DeviceGuardImplInterface.cpp +++ b/c10/core/impl/DeviceGuardImplInterface.cpp @@ -21,11 +21,6 @@ DeviceGuardImplRegistrar::DeviceGuardImplRegistrar( registerDeviceGuard(type, impl); } -namespace { -thread_local std::unique_ptr tls_fake_device_guard = - nullptr; -} // namespace - void ensureCUDADeviceGuardSet() { constexpr auto cuda_idx = static_cast(DeviceType::CUDA); @@ -38,8 +33,8 @@ void ensureCUDADeviceGuardSet() { // In following cases, we override CUDA guard interface with a no-op // device guard. When p->deviceCount() == 0, cuda build is enabled, but no // cuda devices available. - tls_fake_device_guard = std::make_unique>(); - device_guard_impl_registry[cuda_idx].store(tls_fake_device_guard.get()); + static FakeGuardImpl fake_cuda_guard; + device_guard_impl_registry[cuda_idx].store(&fake_cuda_guard); } } diff --git a/c10/core/impl/DeviceGuardImplInterface.h b/c10/core/impl/DeviceGuardImplInterface.h index 01ef32163bb5c..ef7a628835d97 100644 --- a/c10/core/impl/DeviceGuardImplInterface.h +++ b/c10/core/impl/DeviceGuardImplInterface.h @@ -228,6 +228,13 @@ struct C10_API DeviceGuardImplInterface { TORCH_CHECK(false, "Backend doesn't support synchronizing streams."); } + /** + * Return true if this stream is currently recording work for graph capture. + */ + virtual bool isStreamCapturing(const Stream& /*stream*/) const { + TORCH_CHECK(false, "Backend doesn't support stream capture query."); + } + /** * Wait (by blocking the calling thread) until all the work previously * recorded on the event has completed running on the device. diff --git a/c10/core/impl/LocalDispatchKeySet.cpp b/c10/core/impl/LocalDispatchKeySet.cpp index 0b3d87fce2410..c8277f19acc7c 100644 --- a/c10/core/impl/LocalDispatchKeySet.cpp +++ b/c10/core/impl/LocalDispatchKeySet.cpp @@ -46,29 +46,21 @@ void _force_tls_local_dispatch_key_set(LocalDispatchKeySet key_set) { // RAII API IncludeDispatchKeyGuard::IncludeDispatchKeyGuard(DispatchKeySet include) - : tls_(&raw_local_dispatch_key_set), include_(include - tls_->included()) { - if (!include_.empty()) { - tls_->set_included(tls_->included() | include_); - } + : tls_(&raw_local_dispatch_key_set), saved_state_(tls_->included()) { + tls_->set_included(saved_state_ | include); } IncludeDispatchKeyGuard::~IncludeDispatchKeyGuard() { - if (!include_.empty()) { - tls_->set_included(tls_->included() - include_); - } + tls_->set_included(saved_state_); } ExcludeDispatchKeyGuard::ExcludeDispatchKeyGuard(DispatchKeySet exclude) - : tls_(&raw_local_dispatch_key_set), exclude_(exclude - tls_->excluded()) { - if (!exclude_.empty()) { - tls_->set_excluded(tls_->excluded() | exclude_); - } + : tls_(&raw_local_dispatch_key_set), saved_state_(tls_->excluded()) { + tls_->set_excluded(saved_state_ | exclude); } ExcludeDispatchKeyGuard::~ExcludeDispatchKeyGuard() { - if (!exclude_.empty()) { - tls_->set_excluded(tls_->excluded() - exclude_); - } + tls_->set_excluded(saved_state_); } // Non-RAII API diff --git a/c10/core/impl/LocalDispatchKeySet.h b/c10/core/impl/LocalDispatchKeySet.h index bba089bb2ad11..bd8fe811b2c06 100644 --- a/c10/core/impl/LocalDispatchKeySet.h +++ b/c10/core/impl/LocalDispatchKeySet.h @@ -94,7 +94,7 @@ class C10_API IncludeDispatchKeyGuard { // A little micro-optimization to save us from tls_get_addr call // on destruction PODLocalDispatchKeySet* tls_; - DispatchKeySet include_; + DispatchKeySet saved_state_; }; class C10_API ExcludeDispatchKeyGuard { @@ -112,7 +112,7 @@ class C10_API ExcludeDispatchKeyGuard { // A little micro-optimization to save us from tls_get_addr call // on destruction PODLocalDispatchKeySet* tls_; - DispatchKeySet exclude_; + DispatchKeySet saved_state_; }; struct C10_API ForceDispatchKeyGuard { diff --git a/c10/core/impl/VirtualGuardImpl.h b/c10/core/impl/VirtualGuardImpl.h index 6c86a71676e55..2444e4e03b3ee 100644 --- a/c10/core/impl/VirtualGuardImpl.h +++ b/c10/core/impl/VirtualGuardImpl.h @@ -89,6 +89,9 @@ class VirtualGuardImpl final : public DeviceGuardImplInterface { void synchronizeStream(const Stream& stream) const override { impl_->synchronizeStream(stream); } + bool isStreamCapturing(const Stream& stream) const override { + return impl_->isStreamCapturing(stream); + } void recordDataPtrOnStream(const c10::DataPtr& data_ptr, const Stream& stream) const override { diff --git a/c10/cuda/CUDAAllocatorConfig.cpp b/c10/cuda/CUDAAllocatorConfig.cpp index 5414d838cd8c4..f6fb7c9d6d9a8 100644 --- a/c10/cuda/CUDAAllocatorConfig.cpp +++ b/c10/cuda/CUDAAllocatorConfig.cpp @@ -109,6 +109,12 @@ void CUDAAllocatorConfig::parseArgs(const std::string& env) { } else if (key == "per_process_memory_fraction") { i = parsePerProcessMemoryFraction(tokenizer, i); used_native_specific_option = true; + } else if (key == "pinned_free_catch_all") { + i = parsePinnedFreeCatchAll(tokenizer, i); + used_native_specific_option = true; + } else if (key == "throw_on_cudamalloc_oom") { + i = parseThrowOnCudaMallocOom(tokenizer, i); + used_native_specific_option = true; } else { const auto& keys = c10::CachingAllocator::AcceleratorAllocatorConfig::getKeys(); @@ -192,6 +198,25 @@ size_t CUDAAllocatorConfig::parsePinnedReserveSegmentSize( return i; } +size_t CUDAAllocatorConfig::parsePinnedFreeCatchAll( + const c10::CachingAllocator::ConfigTokenizer& tokenizer, + size_t i) { + tokenizer.checkToken(++i, ":"); + m_pinned_free_catch_all = tokenizer.toBool(++i); + return i; +} + +size_t CUDAAllocatorConfig::parseThrowOnCudaMallocOom( + const c10::CachingAllocator::ConfigTokenizer& tokenizer, + size_t i) { + // Format: throw_on_cudamalloc_oom:true or throw_on_cudamalloc_oom:false + // When enabled, throws OOM error before calling cudaMalloc if the allocation + // would likely fail due to insufficient memory. + tokenizer.checkToken(++i, ":"); + m_throw_on_cudamalloc_oom = tokenizer.toBool(++i); + return i; +} + REGISTER_ALLOCATOR_CONFIG_PARSE_HOOK(CUDAAllocatorConfig) } // namespace c10::cuda::CUDACachingAllocator diff --git a/c10/cuda/CUDAAllocatorConfig.h b/c10/cuda/CUDAAllocatorConfig.h index 4e6097a406bc2..928384a874394 100644 --- a/c10/cuda/CUDAAllocatorConfig.h +++ b/c10/cuda/CUDAAllocatorConfig.h @@ -34,7 +34,8 @@ class C10_CUDA_API CUDAAllocatorConfig { static bool expandable_segments() { bool enabled = c10::CachingAllocator::AcceleratorAllocatorConfig:: use_expandable_segments(); -#ifndef PYTORCH_C10_DRIVER_API_SUPPORTED +#if !defined(PYTORCH_C10_DRIVER_API_SUPPORTED) && \ + (!defined(USE_ROCM) || (ROCM_VERSION < 70000)) if (enabled) { TORCH_WARN_ONCE("expandable_segments not supported on this platform") } @@ -65,6 +66,13 @@ class C10_CUDA_API CUDAAllocatorConfig { return instance().m_per_process_memory_fraction; } + // When enabled, throws OOM error before calling cudaMalloc if the allocation + // would likely fail due to insufficient memory. This provides early failure + // with clear error messages instead of letting cudaMalloc fail. + static bool throw_on_cudamalloc_oom() { + return instance().m_throw_on_cudamalloc_oom; + } + /** Pinned memory allocator settings */ static bool pinned_use_cuda_host_register() { return instance().m_pinned_use_cuda_host_register; @@ -92,6 +100,10 @@ class C10_CUDA_API CUDAAllocatorConfig { return 128; } + static bool pinned_free_catch_all() { + return instance().m_pinned_free_catch_all; + } + C10_DEPRECATED_MESSAGE( "c10::cuda::CUDACachingAllocator::CUDAAllocatorConfig::roundup_power2_divisions() is deprecated. Please use c10::CachingAllocator::AcceleratorAllocatorConfig::roundup_power2_divisions() instead.") static size_t roundup_power2_divisions(size_t size) { @@ -157,7 +169,9 @@ class C10_CUDA_API CUDAAllocatorConfig { "graph_capture_record_stream_reuse", "pinned_reserve_segment_size_mb", "pinned_num_register_threads", - "per_process_memory_fraction"}; + "per_process_memory_fraction", + "pinned_free_catch_all", + "throw_on_cudamalloc_oom"}; return keys; } @@ -185,6 +199,12 @@ class C10_CUDA_API CUDAAllocatorConfig { double parsePerProcessMemoryFraction( const c10::CachingAllocator::ConfigTokenizer& tokenizer, size_t i); + size_t parsePinnedFreeCatchAll( + const c10::CachingAllocator::ConfigTokenizer& tokenizer, + size_t i); + size_t parseThrowOnCudaMallocOom( + const c10::CachingAllocator::ConfigTokenizer& tokenizer, + size_t i); std::atomic m_pinned_num_register_threads{1}; std::atomic m_pinned_reserve_segment_size_mb{0}; @@ -198,6 +218,10 @@ class C10_CUDA_API CUDAAllocatorConfig { std::atomic m_pinned_use_cuda_host_register{false}; std::atomic m_graph_capture_record_stream_reuse{false}; std::atomic m_per_process_memory_fraction{1.0}; + std::atomic m_pinned_free_catch_all{false}; + // When true, throw OOM error before calling cudaMalloc if allocation would + // fail + std::atomic m_throw_on_cudamalloc_oom{false}; }; // Keep this for backwards compatibility diff --git a/c10/cuda/CUDACachingAllocator.cpp b/c10/cuda/CUDACachingAllocator.cpp index eb628b51968a4..774c79102b8d6 100644 --- a/c10/cuda/CUDACachingAllocator.cpp +++ b/c10/cuda/CUDACachingAllocator.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -17,16 +18,23 @@ #include #include -#if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED) +#if defined(PYTORCH_C10_DRIVER_API_SUPPORTED) || defined(USE_ROCM) +#if defined(PYTORCH_C10_DRIVER_API_SUPPORTED) #include +#endif +#ifndef _WIN32 #include #include #include +#else +#include +#endif #endif #include #include #include +#include #include #include #include @@ -133,7 +141,7 @@ namespace Native { */ // counter to track order for Mempool Registration -thread_local int32_t registration_counter_global = -1; +std::atomic registration_counter_global{-1}; static char SHAREABLE_HANDLE_VERSION = 2; enum ShareableHandleType : char { @@ -158,12 +166,12 @@ void decrease_stat_array( struct Block; struct PrivatePool; typedef bool (*Comparison)(const Block*, const Block*); -static bool BlockComparatorSize(const Block* a, const Block* b); +static bool BlockComparatorRegistrationCounter(const Block* a, const Block* b); static bool BlockComparatorAddress(const Block* a, const Block* b); struct BlockPool { BlockPool(bool small, PrivatePool* private_pool = nullptr) - : blocks(BlockComparatorSize), + : blocks(BlockComparatorRegistrationCounter), unmapped(BlockComparatorAddress), is_small(small), owner_PrivatePool(private_pool) {} @@ -226,14 +234,16 @@ struct Block { requested_size(0), pool(pool), ptr(ptr) { - registration_counter = ++registration_counter_global; + registration_counter = + registration_counter_global.fetch_add(1, std::memory_order_relaxed) + 1; } // constructor for search key + // Use the default value for registration_counter and not modify + // registration_counter_global, because the search key is just a + // dummy placeholder. Block(c10::DeviceIndex device, cudaStream_t stream, size_t size) - : device(device), stream(stream), size(size), requested_size(0) { - registration_counter = ++registration_counter_global; - } + : device(device), stream(stream), size(size), requested_size(0) {} size_t gc_count() { TORCH_INTERNAL_ASSERT(pool); @@ -269,7 +279,8 @@ struct SegmentRange { SegmentRange(void* p, size_t s) : ptr(static_cast(p)), size(s) {} }; -#if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED) +#if !defined(USE_ROCM) && defined(PYTORCH_C10_DRIVER_API_SUPPORTED) || \ + defined(USE_ROCM) /* Note [Expandable Segments] @@ -383,8 +394,13 @@ struct ExpandableSegment { // This allows for some cases where we have to unmap pages earlier in the // segment to put them at the end. max_handles_ = numSegments(prop.totalGlobalMem + prop.totalGlobalMem / 8); +#ifdef USE_ROCM + C10_CUDA_CHECK(hipMemAddressReserve( + &ptr_, segment_size_ * max_handles_, 0ULL, 0, 0ULL)); +#else C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemAddressReserve_( &ptr_, segment_size_ * max_handles_, 0ULL, 0, 0ULL)); +#endif } ExpandableSegment(const ExpandableSegment&) = delete; ExpandableSegment(ExpandableSegment&&) = delete; @@ -408,12 +424,14 @@ struct ExpandableSegment { // if it fails, use posix file handle if (CUDAAllocatorConfig::expandable_segments_handle_type() == Expandable_Segments_Handle_Type::UNSPECIFIED) { +#ifndef USE_ROCM CUDAAllocatorConfig::set_expandable_segments_handle_type( Expandable_Segments_Handle_Type::FABRIC_HANDLE); auto output = map(range); if (output.ptr != nullptr) { return output; } +#endif // if fabric handle is not supported, use posix file handle. CUDAAllocatorConfig::set_expandable_segments_handle_type( Expandable_Segments_Handle_Type::POSIX_FD); @@ -445,33 +463,60 @@ struct ExpandableSegment { if (enable_ipc_handles) { if (CUDAAllocatorConfig::expandable_segments_handle_type() != Expandable_Segments_Handle_Type::FABRIC_HANDLE) { +#ifdef USE_ROCM + prop.requestedHandleType = hipMemHandleTypePosixFileDescriptor; +#else prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; +#endif } else { +#ifndef USE_ROCM prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; +#endif } } int flag = 0; +#ifndef USE_ROCM C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuDeviceGetAttribute_( &flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, device_)); +#endif if (flag) prop.allocFlags.gpuDirectRDMACapable = 1; prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; // NOLINTNEXTLINE(bugprone-signed-char-misuse) prop.location.id = static_cast(device_); +#ifdef USE_ROCM + auto status = hipMemCreate(&handle, segment_size_, &prop, 0); +#else auto status = DriverAPI::get()->cuMemCreate_(&handle, segment_size_, &prop, 0); +#endif if (status != CUDA_SUCCESS) { if (status == CUDA_ERROR_OUT_OF_MEMORY) { +#ifdef USE_ROCM + // hipMemCreate above returned hipErrorOutOfMemory and treated it + // like a sticky runtime error. Which means we need to clear it. + // Unlike the corresponding CUDA Driver API. + (void)hipGetLastError(); +#endif for (auto j : c10::irange(begin, i)) { // NOLINTNEXTLINE(bugprone-unchecked-optional-access) auto h = handles_.at(j).value(); handles_.at(j) = std::nullopt; +#ifdef USE_ROCM + C10_CUDA_CHECK(hipMemRelease(h.handle)); +#else C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemRelease_(h.handle)); +#endif } trimHandles(); return rangeFromHandles(begin, begin); +#ifdef USE_ROCM + } else { + C10_CUDA_CHECK(status); + } +#else } else if ( CUDAAllocatorConfig::expandable_segments_handle_type() == Expandable_Segments_Handle_Type::FABRIC_HANDLE) { @@ -487,6 +532,7 @@ struct ExpandableSegment { } else { C10_CUDA_DRIVER_CHECK(status); } +#endif } handles_.at(i) = Handle{handle, std::nullopt}; } @@ -522,7 +568,11 @@ struct ExpandableSegment { // thereby ensuring that the handle can be correctly matched in // ipcMemHandle_to_devptr. ShareHeader header{}; +#ifdef _WIN32 + header.pid = _getpid(); +#else header.pid = getpid(); +#endif header.segment_size = segment_size_; header.num_handles = end - begin; @@ -534,8 +584,13 @@ struct ExpandableSegment { Expandable_Segments_Handle_Type::FABRIC_HANDLE) { if (!handle.shareable_handle) { int fd = 0; +#ifdef USE_ROCM + C10_CUDA_CHECK(hipMemExportToShareableHandle( + &fd, handle.handle, hipMemHandleTypePosixFileDescriptor, 0)); +#else C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemExportToShareableHandle_( &fd, handle.handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0)); +#endif handle.shareable_handle = fd; LOG(INFO) << "use posix fd to share expandable segments."; } @@ -546,6 +601,10 @@ struct ExpandableSegment { reinterpret_cast(&*handle.shareable_handle), sizeof(int)); } else { +#ifdef USE_ROCM + TORCH_INTERNAL_ASSERT( + false, "expandable segment with fabric handle not supported"); +#else if (!handle.shareable_handle) { CUmemFabricHandle fabric_handle; C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemExportToShareableHandle_( @@ -559,6 +618,7 @@ struct ExpandableSegment { buf.write( reinterpret_cast(&*handle.shareable_handle), sizeof(CUmemFabricHandle)); +#endif } } return rangeFromHandles(begin, end); @@ -574,14 +634,20 @@ struct ExpandableSegment { device, std::nullopt, header.segment_size, std::move(peers)); // older build setups (e.g. multiwheels) do not have this syscall, added 2020 // but the kernel on the system might still support it. +#ifndef _WIN32 #ifndef SYS_pidfd_open #define SYS_pidfd_open 434 #endif #ifndef SYS_pidfd_getfd #define SYS_pidfd_getfd 438 #endif +#endif // !_WIN32 if (CUDAAllocatorConfig::expandable_segments_handle_type() != Expandable_Segments_Handle_Type::FABRIC_HANDLE) { +#ifdef _WIN32 + TORCH_CHECK( + false, "IPC expandable segments are not supported on Windows"); +#else auto pidfd = syscall(SYS_pidfd_open, header.pid, 0); TORCH_CHECK( pidfd != -1 || errno != ENOSYS, @@ -597,9 +663,13 @@ struct ExpandableSegment { auto err = errno; close(static_cast(pidfd)); for (auto& h : segment->handles_) { +#ifdef USE_ROCM + C10_CUDA_CHECK(hipMemRelease(h.value().handle)); +#else C10_CUDA_DRIVER_CHECK( // NOLINTNEXTLINE(bugprone-unchecked-optional-access) DriverAPI::get()->cuMemRelease_(h.value().handle)); +#endif h = std::nullopt; } TORCH_CHECK( @@ -609,31 +679,56 @@ struct ExpandableSegment { TORCH_CHECK(false, "pidfd_getfd: ", c10::utils::str_error(err)); } CUmemGenericAllocationHandle handle = 0; - C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemImportFromShareableHandle_( - &handle, - // NOLINTNEXTLINE(performance-no-int-to-ptr) - (void*)(uintptr_t)myfd, - CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR)); +#ifdef USE_ROCM +#if ROCM_VERSION >= 70100 + void* myfd_handle = + reinterpret_cast(static_cast(myfd)); +#else + void* myfd_handle = (void*)(uintptr_t)&myfd; +#endif + C10_CUDA_CHECK(hipMemImportFromShareableHandle( + &handle, myfd_handle, hipMemHandleTypePosixFileDescriptor)); +#else + C10_CUDA_DRIVER_CHECK_MSG( + DriverAPI::get()->cuMemImportFromShareableHandle_( + &handle, + // NOLINTNEXTLINE(performance-no-int-to-ptr) + (void*)(uintptr_t)myfd, + CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR), + " fabric_info: {", + get_nvml_fabric_info(device), + "}"); +#endif LOG(INFO) << "use posix fd to import expandable segments."; close(static_cast(myfd)); segment->handles_.emplace_back(Handle{handle, std::nullopt}); } close(static_cast(pidfd)); +#endif // !_WIN32 } else { +#ifdef USE_ROCM + TORCH_INTERNAL_ASSERT( + false, "expandable segment with fabric handle not supported"); +#else for (auto i : c10::irange(header.num_handles)) { (void)i; CUmemFabricHandle fabric_handle; buf.read( reinterpret_cast(&fabric_handle), sizeof(CUmemFabricHandle)); CUmemGenericAllocationHandle handle = 0; - C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemImportFromShareableHandle_( - &handle, - // NOLINTNEXTLINE(performance-no-int-to-ptr) - (void*)&fabric_handle, - CU_MEM_HANDLE_TYPE_FABRIC)); + C10_CUDA_DRIVER_CHECK_MSG( + DriverAPI::get()->cuMemImportFromShareableHandle_( + &handle, + // NOLINTNEXTLINE(performance-no-int-to-ptr) + (void*)&fabric_handle, + CU_MEM_HANDLE_TYPE_FABRIC), + " fabric_info: {", + get_nvml_fabric_info(device), + "}"); LOG(INFO) << "use fabric handle to import expandable segments."; segment->handles_.emplace_back(Handle{handle, std::nullopt}); } +#endif } segment->mapAndSetAccess(0, header.num_handles); return segment; @@ -669,23 +764,55 @@ struct ExpandableSegment { ~ExpandableSegment() { forEachAllocatedRange( [&](size_t begin, size_t end) { unmapHandles(begin, end); }); +#ifdef USE_ROCM + C10_CUDA_CHECK(hipMemAddressFree(ptr_, segment_size_ * max_handles_)); +#else C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemAddressFree_( ptr_, segment_size_ * max_handles_)); +#endif } private: void setAccess(c10::DeviceIndex device, size_t begin, size_t end) { - CUmemAccessDesc desc; - desc.location.type = CU_MEM_LOCATION_TYPE_DEVICE; +#if defined(USE_ROCM) && (ROCM_VERSION >= 70200) + constexpr int num_desc = 2; + CUmemAccessDesc desc[num_desc]; + desc[1].location.type = CU_MEM_LOCATION_TYPE_HOST; + desc[1].location.id = 0; // ignored + desc[1].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#else + constexpr int num_desc = 1; + CUmemAccessDesc desc[num_desc]; +#endif + desc[0].location.type = CU_MEM_LOCATION_TYPE_DEVICE; // NOLINTNEXTLINE(bugprone-signed-char-misuse) - desc.location.id = static_cast(device); - desc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + desc[0].location.id = static_cast(device); + desc[0].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; +#ifdef USE_ROCM + C10_CUDA_CHECK(hipMemSetAccess( + ptr() + begin * segment_size_, + (end - begin) * segment_size_, + &desc[0], + num_desc)); +#else C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemSetAccess_( - ptr_ + begin * segment_size_, (end - begin) * segment_size_, &desc, 1)); + ptr_ + begin * segment_size_, + (end - begin) * segment_size_, + &desc[0], + num_desc)); +#endif } void mapAndSetAccess(size_t begin, size_t end) { for (auto i : c10::irange(begin, end)) { +#ifdef USE_ROCM + C10_CUDA_CHECK(hipMemMap( + ptr() + i * segment_size_, + segment_size_, + 0, + handles_.at(i).value().handle, + 0ULL)); +#else C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemMap_( ptr_ + i * segment_size_, segment_size_, @@ -693,6 +820,7 @@ struct ExpandableSegment { // NOLINTNEXTLINE(bugprone-unchecked-optional-access) handles_.at(i).value().handle, 0ULL)); +#endif } mapped_size_ += (end - begin) * segment_size_; setAccess(device_, begin, end); @@ -719,12 +847,22 @@ struct ExpandableSegment { // NOLINTNEXTLINE(bugprone-unchecked-optional-access) Handle h = handles_.at(i).value(); handles_.at(i) = std::nullopt; +#ifdef USE_ROCM + C10_CUDA_CHECK(hipMemUnmap(ptr() + segment_size_ * i, segment_size_)); +#else C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemUnmap_( ptr_ + segment_size_ * i, segment_size_)); +#endif if (h.shareable_handle) { +#ifndef _WIN32 close(std::get(*h.shareable_handle)); +#endif } +#ifdef USE_ROCM + C10_CUDA_CHECK(hipMemRelease(h.handle)); +#else C10_CUDA_DRIVER_CHECK(DriverAPI::get()->cuMemRelease_(h.handle)); +#endif } trimHandles(); } @@ -772,7 +910,11 @@ struct ExpandableSegment { std::optional> shareable_handle; }; struct ShareHeader { +#ifdef _WIN32 + int pid; +#else pid_t pid; +#endif size_t segment_size; size_t num_handles; }; @@ -866,15 +1008,16 @@ struct RestoreResult { std::vector allocations_created; }; -bool BlockComparatorSize(const Block* a, const Block* b) { +bool BlockComparatorRegistrationCounter(const Block* a, const Block* b) { if (a->stream != b->stream) { return (uintptr_t)a->stream < (uintptr_t)b->stream; } if (a->size != b->size) { return a->size < b->size; } - return (uintptr_t)a->ptr < (uintptr_t)b->ptr; + return a->registration_counter < b->registration_counter; } + bool BlockComparatorAddress(const Block* a, const Block* b) { if (a->stream != b->stream) { return (uintptr_t)a->stream < (uintptr_t)b->stream; @@ -882,6 +1025,14 @@ bool BlockComparatorAddress(const Block* a, const Block* b) { return (uintptr_t)a->ptr < (uintptr_t)b->ptr; } +// Info about OOM rejection, used to defer observer callbacks outside of lock +struct OomRejectionInfo { + bool rejected{false}; + size_t alloc_size{0}; + size_t total_allocated{0}; + size_t device_total{0}; +}; + struct AllocParams { AllocParams( c10::DeviceIndex device, @@ -912,6 +1063,7 @@ struct AllocParams { Block* block{nullptr}; StatTypes stat_types = {false}; cudaError_t err{cudaSuccess}; + OomRejectionInfo oom_rejection_info; }; // Note: cudaEventCreate when concurrently invoked from multiple threads can be @@ -1275,6 +1427,7 @@ class DeviceCachingAllocator { // XXX - maybe we should generalize and have multiple events std::vector oom_observers_; + std::vector oom_rejection_observers_; std::vector trace_trackers_; @@ -1383,6 +1536,10 @@ class DeviceCachingAllocator { oom_observers_.emplace_back(std::move(observer)); } + void attachOomRejectionObserver(OomRejectionObserver observer) { + oom_rejection_observers_.emplace_back(std::move(observer)); + } + void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) { std::unique_lock lock(mutex); trace_trackers_.emplace_back(std::move(tracker)); @@ -1463,19 +1620,28 @@ class DeviceCachingAllocator { // cudaMalloc. So far this function has not modified allocator state, but // keep in mind that any observed allocator state may change across calls // to alloc_block since it may release the lock. - block_found = alloc_block(params, false, context, lock) - // Try to use memory pools that have opted in as overflow before - // expensive memory freeing operations. - || try_mempool_fallback( - params, size, stream, device_id, alloc_size, stats) - // Free enough available cached blocks to satisfy alloc and retry - // alloc. - || (release_available_cached_blocks(params, context) && - alloc_block(params, false, context, lock)) - // Free all non-split cached blocks and retry alloc. - || (C10_LIKELY(captures_underway.empty()) && - release_cached_blocks(context, {0, 0}) && - alloc_block(params, true, context, lock)); + block_found = alloc_block(params, false, context, lock); + + // If allocation was rejected by OOM policy, skip retry chain and fail + // immediately + if (!block_found && params.oom_rejection_info.rejected) { + // Skip retry chain - will be handled below in the !block_found path + } else if (!block_found) { + // Normal retry chain: try various strategies to free memory and retry + block_found = + // Try to use memory pools that have opted in as overflow before + // expensive memory freeing operations. + try_mempool_fallback( + params, size, stream, device_id, alloc_size, stats) + // Free enough available cached blocks to satisfy alloc and retry + // alloc. + || (release_available_cached_blocks(params, context) && + alloc_block(params, false, context, lock)) + // Free all non-split cached blocks and retry alloc. + || (C10_LIKELY(captures_underway.empty()) && + release_cached_blocks(context, {0, 0}) && + alloc_block(params, true, context, lock)); + } } if (!block_found) { @@ -1483,6 +1649,43 @@ class DeviceCachingAllocator { // alloc_block should have thrown an exception already. TORCH_INTERNAL_ASSERT(params.err == cudaErrorMemoryAllocation); + // Handle OOM rejection separately - return nullptr instead of throwing + // This allows callers to handle rejection gracefully without crashing + if (params.oom_rejection_info.rejected) { + if (!oom_rejection_observers_.empty()) { + auto observers_local = oom_rejection_observers_; + auto rejection_info = params.oom_rejection_info; + auto device = device_id; + + // Release lock before dispatching observers + lock.unlock(); + + // Dispatch observers asynchronously to minimize latency impact on + // rejection path + std::thread([observers_local = std::move(observers_local), + rejection_info, + device]() { + try { + for (const auto& observer : observers_local) { + observer( + device, + rejection_info.alloc_size, + rejection_info.total_allocated, + rejection_info.device_total); + } + } catch (const std::exception& e) { + LOG(ERROR) << "Exception in OOM rejection observer: " << e.what(); + } catch (...) { + LOG(ERROR) << "Unknown exception in OOM rejection observer"; + } + }).detach(); + } else { + lock.unlock(); + } + + return nullptr; + } + size_t device_free = 0; size_t device_total = 0; C10_CUDA_CHECK(cudaMemGetInfo(&device_free, &device_total)); @@ -2182,6 +2385,7 @@ class DeviceCachingAllocator { stats.num_sync_all_streams = 0; stats.num_device_alloc = 0; stats.num_device_free = 0; + stats.num_oom_rejections = 0; stats.oversize_allocations.reset_accumulated(); stats.oversize_segments.reset_accumulated(); } @@ -3340,8 +3544,37 @@ class DeviceCachingAllocator { total_allocated_memory + size > allowed_memory_maximum.value()) { p.err = cudaErrorMemoryAllocation; return false; - // Temporarily disable checkpointing & cudagraphs internally - } else if ( + } + + // When throw_on_cudamalloc_oom is enabled, check + // per_process_memory_fraction to reject allocations that would exceed the + // configured limit. This prevents fatal HSA/driver aborts by throwing + // OutOfMemoryError instead. + if (CUDAAllocatorConfig::throw_on_cudamalloc_oom()) { + size_t device_total = static_cast(device_prop.totalGlobalMem); + size_t max_allowed = static_cast( + CUDAAllocatorConfig::per_process_memory_fraction() * + static_cast(device_total)); + if (total_allocated_memory + size > max_allowed) { + stats.num_oom_rejections++; + p.oom_rejection_info = { + true, size, total_allocated_memory, device_total}; + C10_LOG_EVERY_MS(WARNING, 60000) + << "Preemptively rejecting allocation: requested " + << format_size(size) << " but total_allocated_memory (" + << format_size(total_allocated_memory) + << ") + requested size would exceed memory limit (" + << format_size(max_allowed) << " = " + << CUDAAllocatorConfig::per_process_memory_fraction() * 100 + << "% of " << format_size(device_total) + << "). Set throw_on_cudamalloc_oom:false to disable."; + p.err = cudaErrorMemoryAllocation; + return false; + } + } + + if ( + // Temporarily disable checkpointing & cudagraphs internally p.is_expandable_segments_active && !(in_fbcode && p.pool->owner_PrivatePool)) { p.block = try_allocate_expandable_block( @@ -3883,7 +4116,7 @@ class DeviceCachingAllocator { if (record_history) { // Skip if action is in the skip_actions set - bool should_skip = skip_actions_list.count(action) > 0; + bool should_skip = skip_actions_list.contains(action); if (!should_skip) { alloc_buffer.insertEntries(te); } @@ -4013,6 +4246,18 @@ class NativeCachingAllocator : public CUDAAllocator { device, ": did you call init?"); Block* block = device_allocator[device]->malloc(size, stream); + // block can be nullptr if allocation was rejected by + // throw_on_cudamalloc_oom + // + per_process_memory_fraction policy. Throw OutOfMemoryError so callers + // with OOM error handlers can catch it gracefully. + if (C10_UNLIKELY(!block)) { + TORCH_CHECK_WITH( + OutOfMemoryError, + false, + "CUDA out of memory. Allocation was preemptively rejected because " + "it would exceed per_process_memory_fraction limit. Requested size: ", + format_size(size)); + } add_allocated_block(block); *devPtr = block->ptr; const c10::impl::PyInterpreter* interp = c10::impl::GPUTrace::get_trace(); @@ -4164,6 +4409,12 @@ class NativeCachingAllocator : public CUDAAllocator { } } + void attachOomRejectionObserver(OomRejectionObserver observer) override { + for (auto& allocator : device_allocator) { + allocator->attachOomRejectionObserver(observer); + } + } + void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) override { for (auto& allocator : device_allocator) { allocator->attachAllocatorTraceTracker(tracker); diff --git a/c10/cuda/CUDACachingAllocator.h b/c10/cuda/CUDACachingAllocator.h index 9cd6b4ee128b6..c2d502976dda6 100644 --- a/c10/cuda/CUDACachingAllocator.h +++ b/c10/cuda/CUDACachingAllocator.h @@ -95,14 +95,26 @@ using OutOfMemoryObserver = std::function; +// Observer called when an allocation is preemptively rejected due to +// throw_on_cudamalloc_oom policy. Parameters: +// - device: GPU device index +// - alloc_size: size of the rejected allocation request +// - total_allocated: total memory allocated before the request +// - device_total: total GPU memory +using OomRejectionObserver = std::function; + struct ShareableHandle { ptrdiff_t offset; std::string handle; }; struct StreamSegmentSize { - StreamSegmentSize(cudaStream_t s, bool small, size_t sz) - : stream(s), is_small_pool(small), total_size(sz) {} + StreamSegmentSize(cudaStream_t s, bool small_, size_t sz) + : stream(s), is_small_pool(small_), total_size(sz) {} cudaStream_t stream; bool is_small_pool; size_t total_size; @@ -216,6 +228,7 @@ class CUDAAllocator : public DeviceAllocator { return ""; } virtual void attachOutOfMemoryObserver(OutOfMemoryObserver observer) = 0; + virtual void attachOomRejectionObserver(OomRejectionObserver observer) = 0; // Attached AllocatorTraceTracker callbacks will be called while the // per-device allocator lock is held. Any additional locks taken from within @@ -421,6 +434,10 @@ inline void attachOutOfMemoryObserver(OutOfMemoryObserver observer) { get()->attachOutOfMemoryObserver(std::move(observer)); } +inline void attachOomRejectionObserver(OomRejectionObserver observer) { + get()->attachOomRejectionObserver(std::move(observer)); +} + inline void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) { get()->attachAllocatorTraceTracker(std::move(tracker)); } diff --git a/c10/cuda/CUDAGraphsC10Utils.h b/c10/cuda/CUDAGraphsC10Utils.h index 936875fd71d5c..3709aa64c7845 100644 --- a/c10/cuda/CUDAGraphsC10Utils.h +++ b/c10/cuda/CUDAGraphsC10Utils.h @@ -1,8 +1,10 @@ #pragma once +#include #include + #include -#include +#include // CUDA Graphs utils used by c10 and aten. // aten/cuda/CUDAGraphsUtils.cuh adds utils used by aten only. @@ -67,10 +69,41 @@ inline std::ostream& operator<<(std::ostream& os, CaptureStatus status) { // Use this version where you're sure a CUDA context exists already. inline CaptureStatus currentStreamCaptureStatusMayInitCtx() { - cudaStreamCaptureStatus is_capturing{cudaStreamCaptureStatusNone}; + cudaStreamCaptureStatus status{cudaStreamCaptureStatusNone}; C10_CUDA_CHECK( - cudaStreamIsCapturing(c10::cuda::getCurrentCUDAStream(), &is_capturing)); - return CaptureStatus(is_capturing); + cudaStreamIsCapturing(c10::cuda::getCurrentCUDAStream(), &status)); + return CaptureStatus(status); +} + +inline CaptureStatus captureStatusMayInitCtx(cudaStream_t stream) { + cudaStreamCaptureStatus status{cudaStreamCaptureStatusNone}; + C10_CUDA_CHECK(cudaStreamIsCapturing(stream, &status)); + return CaptureStatus(status); +} + +inline bool isStreamCapturingMayInitCtx(cudaStream_t stream) { + return captureStatusMayInitCtx(stream) == CaptureStatus::Active; +} + +inline std::optional currentStreamCaptureIdMayInitCtx() { + cudaStreamCaptureStatus status{}; + CaptureId_t capture_id = 0; + C10_CUDA_CHECK(cudaStreamGetCaptureInfo( + c10::cuda::getCurrentCUDAStream(), &status, &capture_id)); + if (status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive) { + return capture_id; + } + return std::nullopt; +} + +inline std::optional captureIdMayInitCtx(cudaStream_t stream) { + cudaStreamCaptureStatus status{}; + CaptureId_t capture_id = 0; + C10_CUDA_CHECK(cudaStreamGetCaptureInfo(stream, &status, &capture_id)); + if (status == cudaStreamCaptureStatus::cudaStreamCaptureStatusActive) { + return capture_id; + } + return std::nullopt; } } // namespace c10::cuda diff --git a/c10/cuda/CUDAMallocAsyncAllocator.cpp b/c10/cuda/CUDAMallocAsyncAllocator.cpp index cc8bec049fc3f..7f9e78c8be861 100644 --- a/c10/cuda/CUDAMallocAsyncAllocator.cpp +++ b/c10/cuda/CUDAMallocAsyncAllocator.cpp @@ -668,6 +668,13 @@ struct CudaMallocAsyncAllocator : public CUDAAllocator { "If you need it, please file an issue describing your use case."); } + void attachOomRejectionObserver(OomRejectionObserver observer) override { + TORCH_CHECK( + false, + "cudaMallocAsync does not yet support attachOomRejectionObserver. " + "If you need it, please file an issue describing your use case."); + } + void attachAllocatorTraceTracker(AllocatorTraceTracker tracker) override { TORCH_CHECK( false, diff --git a/c10/cuda/CUDAStream.h b/c10/cuda/CUDAStream.h index f27c7c9176631..183232c4d9191 100644 --- a/c10/cuda/CUDAStream.h +++ b/c10/cuda/CUDAStream.h @@ -116,6 +116,13 @@ class C10_CUDA_API CUDAStream { void synchronize() const; + bool is_capturing() const { + DeviceGuard guard{stream_.device()}; + cudaStreamCaptureStatus status{cudaStreamCaptureStatusNone}; + C10_CUDA_CHECK(cudaStreamIsCapturing(stream(), &status)); + return status != cudaStreamCaptureStatusNone; + } + int priority() const { DeviceGuard guard{stream_.device()}; int priority = 0; diff --git a/c10/cuda/PeerToPeerAccess.cpp b/c10/cuda/PeerToPeerAccess.cpp index 526d794bc49d5..cc312bdfaa4cb 100644 --- a/c10/cuda/PeerToPeerAccess.cpp +++ b/c10/cuda/PeerToPeerAccess.cpp @@ -10,12 +10,15 @@ #include #include +#include +#include #include namespace c10::cuda { static std::vector p2pAccessEnabled_; static std::vector fabricAccessEnabled_; +static std::vector fabricCliqueId_; static int64_t num_devices_ = -1; namespace detail { @@ -35,6 +38,8 @@ void init_p2p_access_cache(int64_t num_devices) { } fabricAccessEnabled_.clear(); fabricAccessEnabled_.resize(num_devices, -1); + fabricCliqueId_.clear(); + fabricCliqueId_.resize(num_devices, kCliqueIdNotQueried); } } // namespace detail @@ -180,6 +185,7 @@ bool get_fabric_access(c10::DeviceIndex dev) { fabricInfo.version = nvmlGpuFabricInfo_v2; if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_ == nullptr) { cache = 0; + fabricCliqueId_[dev] = kCliqueIdUnsupported; return false; } TORCH_CHECK( @@ -188,13 +194,17 @@ bool get_fabric_access(c10::DeviceIndex dev) { nvml_device, &fabricInfo)); auto state = fabricInfo.state != NVML_GPU_FABRIC_STATE_NOT_SUPPORTED; if (state) { + fabricCliqueId_[dev] = static_cast(fabricInfo.cliqueId); // now perform the full cycle of allocating - exporting - importing memory state = isFabricSupported(); + } else { + fabricCliqueId_[dev] = kCliqueIdUnsupported; } cache = state ? 1 : 0; return cache; } else { cache = 0; + fabricCliqueId_[dev] = kCliqueIdUnsupported; return false; } #else @@ -203,4 +213,85 @@ bool get_fabric_access(c10::DeviceIndex dev) { #endif } +int get_fabric_clique_id(c10::DeviceIndex dev) { +#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040 && \ + defined(PYTORCH_C10_DRIVER_API_SUPPORTED) + // Ensure cache is populated via get_fabric_access (which does the NVML query + // and stashes clique_id as a side effect). + get_fabric_access(dev); + return fabricCliqueId_[dev]; +#else + (void)dev; + return kCliqueIdUnsupported; +#endif +} + +std::string get_nvml_fabric_info([[maybe_unused]] c10::DeviceIndex dev) { +#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12040 && \ + defined(PYTORCH_C10_DRIVER_API_SUPPORTED) + if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_ == nullptr) { + return "fabric info unsupported (nvmlDeviceGetGpuFabricInfoV not available)"; + } + + auto nvml_device = get_nvml_device(dev); + if (nvml_device == nullptr) { + return "fabric info unknown (failed to get NVML device handle)"; + } + + nvmlGpuFabricInfoV_t info{}; +#ifdef nvmlGpuFabricInfo_v3 + bool has_health_summary = false; + info.version = nvmlGpuFabricInfo_v3; + if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_(nvml_device, &info) == + NVML_SUCCESS) { + has_health_summary = true; + } else +#endif + { + info = {}; + info.version = nvmlGpuFabricInfo_v2; + if (DriverAPI::get()->nvmlDeviceGetGpuFabricInfoV_(nvml_device, &info) != + NVML_SUCCESS) { + return "fabric info unknown (nvmlDeviceGetGpuFabricInfoV failed)"; + } + } + + char uuid_hex[33]; + for (int i = 0; i < 16; ++i) { + snprintf(uuid_hex + i * 2, 3, "%02x", info.clusterUuid[i]); + } + + const char* state_str = "unknown"; + switch (info.state) { + case NVML_GPU_FABRIC_STATE_NOT_SUPPORTED: + state_str = "not_supported"; + break; + case NVML_GPU_FABRIC_STATE_NOT_STARTED: + state_str = "not_started"; + break; + case NVML_GPU_FABRIC_STATE_IN_PROGRESS: + state_str = "in_progress"; + break; + case NVML_GPU_FABRIC_STATE_COMPLETED: + state_str = "completed"; + break; + } + + std::ostringstream oss; + oss << "clique_id=" << info.cliqueId << ", cluster_uuid=" << uuid_hex + << ", state=" << state_str << ", status=" << info.status + << ", health_mask=0x" << std::hex << std::setfill('0') << std::setw(8) + << info.healthMask; +#ifdef nvmlGpuFabricInfo_v3 + if (has_health_summary) { + oss << ", health_summary=" << std::dec + << static_cast(info.healthSummary); + } +#endif + return oss.str(); +#else + return "fabric info unsupported (requires CUDA >= 12.4)"; +#endif +} + } // namespace c10::cuda diff --git a/c10/cuda/PeerToPeerAccess.h b/c10/cuda/PeerToPeerAccess.h index d526843974264..8615f1f335d96 100644 --- a/c10/cuda/PeerToPeerAccess.h +++ b/c10/cuda/PeerToPeerAccess.h @@ -5,6 +5,7 @@ #include #include +#include namespace c10::cuda { @@ -32,4 +33,17 @@ C10_CUDA_API bool get_p2p_access( /// @return true if fabric access is available, false otherwise. C10_CUDA_API bool get_fabric_access(c10::DeviceIndex device); +constexpr int kCliqueIdNotQueried = -2; +constexpr int kCliqueIdUnsupported = -1; + +/// Query the NVLink fabric clique ID for a device. +/// Returns the clique ID (>= 0) if fabric is supported, or kCliqueIdUnsupported +/// if unsupported. +C10_CUDA_API int get_fabric_clique_id(c10::DeviceIndex device); + +/// Returns a formatted string with NVML fabric info (clique_id, cluster_uuid, +/// state, status, health_mask) for the given device. Intended for error +/// diagnostics — only call on failure paths. +C10_CUDA_API std::string get_nvml_fabric_info(c10::DeviceIndex device); + } // namespace c10::cuda diff --git a/c10/cuda/driver_api.h b/c10/cuda/driver_api.h index 970a3f4584117..681ebf49a1452 100644 --- a/c10/cuda/driver_api.h +++ b/c10/cuda/driver_api.h @@ -20,6 +20,23 @@ } \ } while (0) +// clang-format off +#define C10_CUDA_DRIVER_CHECK_MSG(EXPR, ...) \ + do { \ + CUresult __err = EXPR; \ + if (__err != CUDA_SUCCESS) { \ + const char* err_str; \ + CUresult get_error_str_err [[maybe_unused]] = \ + c10::cuda::DriverAPI::get()->cuGetErrorString_(__err, &err_str); \ + if (get_error_str_err != CUDA_SUCCESS) { \ + TORCH_CHECK(false, "CUDA driver error: unknown error", __VA_ARGS__);\ + } else { \ + TORCH_CHECK(false, "CUDA driver error: ", err_str, __VA_ARGS__); \ + } \ + } \ + } while (0) +// clang-format on + #define C10_CUDA_DRIVER_CHECK_GOTO(EXPR, NEXT) \ do { \ CUresult __err = EXPR; \ diff --git a/c10/cuda/impl/CUDAGuardImpl.h b/c10/cuda/impl/CUDAGuardImpl.h index 14a398ee4dc37..c9617271cd498 100644 --- a/c10/cuda/impl/CUDAGuardImpl.h +++ b/c10/cuda/impl/CUDAGuardImpl.h @@ -213,6 +213,11 @@ struct CUDAGuardImpl final : public c10::impl::DeviceGuardImplInterface { cuda_stream.synchronize(); } + bool isStreamCapturing(const Stream& stream) const override { + CUDAStream cuda_stream{stream}; + return cuda_stream.is_capturing(); + } + void synchronizeEvent(void* event) const override { if (!event) return; diff --git a/c10/metal/indexing.h b/c10/metal/indexing.h index 7d930a4665067..ffbec0a3d71d1 100644 --- a/c10/metal/indexing.h +++ b/c10/metal/indexing.h @@ -53,6 +53,25 @@ kernel void unary_dense( output[index] = f(input[index]); } +template +kernel void unary_dense_vec4( + device result_of* output [[buffer(0)]], + constant T* input [[buffer(1)]], + constant uint& numel [[buffer(2)]], + uint index [[thread_position_in_grid]]) { + F f; + uint base = index * 4; + if (base + 4 <= numel) { + using ::metal::vec; + vec val = *(constant vec*)(input + base); + *(device vec, 4>*)(output + base) = { + f(val.x), f(val.y), f(val.z), f(val.w)}; + } else { + for (uint i = base; i < numel; i++) + output[i] = f(input[i]); + } +} + template kernel void unary_strided( device result_of* output [[buffer(0)]], @@ -90,6 +109,18 @@ kernel void unary_strided( constant uint& ndim, \ uint index) +#define REGISTER_UNARY_VEC4_OP(NAME, DTYPE0, DTYPE1) \ + static_assert( \ + ::metal:: \ + is_same_v>, \ + "Output dtype mismatch for unary op " #NAME " and input " #DTYPE0); \ + template [[host_name(#NAME "_dense_vec4_" #DTYPE1 "_" #DTYPE0)]] \ + kernel void ::c10::metal::unary_dense_vec4( \ + device ::c10::metal::result_of * output, \ + constant DTYPE0 * input, \ + constant uint & numel, \ + uint index) + #define DEFINE_UNARY_FLOATING_FUNCTOR(NAME) \ struct NAME##_functor { \ template \ diff --git a/c10/metal/random.h b/c10/metal/random.h index c03d9b8a3149c..647b860b3f22c 100644 --- a/c10/metal/random.h +++ b/c10/metal/random.h @@ -33,7 +33,7 @@ uint4 single_round(uint4 ctr, uint2 key) { constexpr uint kPhiloxSB = 0xCD9E8D57; auto rc0 = mulhilo(kPhiloxSA, ctr.x); auto rc1 = mulhilo(kPhiloxSB, ctr.z); - return uint4(rc1.y ^ ctr.y ^ key.x, rc1.x, rc0.y ^ ctr.w ^ key.y, rc0.x); + return uint4(rc1.x ^ ctr.y ^ key.x, rc1.y, rc0.x ^ ctr.w ^ key.y, rc0.y); } uint4 multiple_rounds(uint4 ctr, uint2 key, uint rounds) { diff --git a/c10/test/CMakeLists.txt b/c10/test/CMakeLists.txt index 27385159976f4..9a576bb996f22 100644 --- a/c10/test/CMakeLists.txt +++ b/c10/test/CMakeLists.txt @@ -18,12 +18,12 @@ if(BUILD_TEST) endforeach() endif() -# ---[ C++17 header compilation test +# ---[ C++17/20 header compilation test if(BUILD_TEST) add_executable(c10_cpp17_header_build_test util/cpp17_header_build_check.cpp) target_link_libraries(c10_cpp17_header_build_test ${C10_LIB} gmock gtest gtest_main) set_target_properties(c10_cpp17_header_build_test PROPERTIES - CXX_STANDARD 17 + CXX_STANDARD 20 CXX_STANDARD_REQUIRED ON ) add_test(NAME c10_cpp17_header_build_test COMMAND $) diff --git a/c10/test/core/DeviceGuard_test.cpp b/c10/test/core/DeviceGuard_test.cpp index a6148adf8af25..6ae4fc1dd6774 100644 --- a/c10/test/core/DeviceGuard_test.cpp +++ b/c10/test/core/DeviceGuard_test.cpp @@ -1,8 +1,12 @@ #include #include +#include #include +#include +#include + using namespace c10; using namespace c10::impl; @@ -39,3 +43,75 @@ TEST(OptionalDeviceGuard, ResetDeviceDifferentDeviceType) { ASSERT_EQ(g.current_device(), Device(DeviceType::HIP, 2)); ASSERT_EQ(g.original_device(), Device(DeviceType::HIP, 0)); } + +// -- ensureCUDADeviceGuardSet ------------------------------------------- + +// Regression test: ensureCUDADeviceGuardSet() used to store a thread-local +// FakeGuardImpl* in the global device_guard_impl_registry. When the owning +// thread exited its TLS was freed, leaving a dangling pointer that the next +// thread to call deviceCount() would dereference (segfault). +// +// The fix is a function-local static, which has program lifetime. We verify +// that the pointer in the registry is still valid (and returns the expected +// deviceCount) after the threads that triggered guard installation have exited. +TEST(EnsureCUDADeviceGuard, NoUseAfterFreeWhenThreadsExit) { + // Simulate "CUDA compiled, no devices visible": a guard that is non-null but + // returns deviceCount() == 0, which is the condition that triggers fake guard + // installation in ensureCUDADeviceGuardSet(). + struct ZeroDeviceGuardImpl final : public DeviceGuardImplInterface { + DeviceType type() const override { + return DeviceType::CUDA; + } + Device exchangeDevice(Device d) const override { + return d; + } + Device getDevice() const override { + return Device(DeviceType::CUDA, 0); + } + void setDevice(Device) const override {} + void uncheckedSetDevice(Device) const noexcept override {} + Stream getStream(Device d) const noexcept override { + return Stream(Stream::UNSAFE, d, 0); + } + Stream exchangeStream(Stream s) const noexcept override { + return s; + } + DeviceIndex deviceCount() const noexcept override { + return 0; + } + void record(void**, const Stream&, const DeviceIndex, const EventFlag) + const override {} + void block(void*, const Stream&) const override {} + bool queryEvent(void*) const override { + return true; + } + void destroyEvent(void*, const DeviceIndex) const noexcept override {} + }; + + constexpr auto cuda_idx = static_cast(DeviceType::CUDA); + const auto* saved = device_guard_impl_registry[cuda_idx].load(); + + static ZeroDeviceGuardImpl zero_impl; + device_guard_impl_registry[cuda_idx].store(&zero_impl); + + // Phase 1: threads call ensureCUDADeviceGuardSet(), detect deviceCount()==0, + // and install a FakeGuardImpl in the global registry. + { + std::vector threads; + for (int i = 0; i < 4; i++) { + threads.emplace_back(ensureCUDADeviceGuardSet); + } + for (auto& t : threads) { + t.join(); + } + } + // The threads' TLS is now destroyed. With the old code the registry now + // holds a dangling pointer; with the fix it holds &fake_cuda_guard (static). + + // Phase 2: the pointer must still be valid and return the expected count. + const auto* p = device_guard_impl_registry[cuda_idx].load(); + ASSERT_NE(p, nullptr); + ASSERT_EQ(p->deviceCount(), kFakeGuardImplMaxDevices); + + device_guard_impl_registry[cuda_idx].store(saved); +} diff --git a/c10/test/core/DispatchKeySet_test.cpp b/c10/test/core/DispatchKeySet_test.cpp index cdbdc150167e0..9066d6100a3db 100644 --- a/c10/test/core/DispatchKeySet_test.cpp +++ b/c10/test/core/DispatchKeySet_test.cpp @@ -404,7 +404,7 @@ TEST(DispatchKeySet, TestBackendComponentToString) { auto k = static_cast(i); auto res = std::string(toString(k)); ASSERT_FALSE(res == "UNKNOWN_BACKEND_BIT"); - ASSERT_FALSE(seen_strings.count(res) > 0); + ASSERT_FALSE(seen_strings.contains(res)); seen_strings.insert(res); } } @@ -439,7 +439,7 @@ TEST(DispatchKeySet, TestFunctionalityDispatchKeyToString) { } else { ASSERT_TRUE(res.find("Unknown") == std::string::npos) << i; } - ASSERT_TRUE(seen_strings.count(res) == 0); + ASSERT_TRUE(!seen_strings.contains(res)); seen_strings.insert(res); } } diff --git a/c10/util/ApproximateClock.h b/c10/util/ApproximateClock.h index 803ea404c9aa2..c1b65b33af5a7 100644 --- a/c10/util/ApproximateClock.h +++ b/c10/util/ApproximateClock.h @@ -29,6 +29,8 @@ #else #undef C10_RDTSC #endif +#elif defined(__aarch64__) && !defined(__CUDACC__) && !defined(__HIPCC__) +#define C10_ARMTSC #endif namespace c10 { @@ -70,6 +72,14 @@ inline time_t getTime(bool allow_monotonic = false) { #endif } +#if defined(C10_ARMTSC) +inline uint64_t getArmApproximateTime() { + uint64_t val; + __asm__ __volatile__("mrs %0, cntvct_el0" : "=r"(val)); + return val; +} +#endif + // We often do not need to capture true wall times. If a fast mechanism such // as TSC is available we can use that instead and convert back to epoch time // during post processing. This greatly reduce the clock's contribution to @@ -81,6 +91,8 @@ inline time_t getTime(bool allow_monotonic = false) { inline auto getApproximateTime() { #if defined(C10_RDTSC) return static_cast(__rdtsc()); +#elif defined(C10_ARMTSC) + return getArmApproximateTime(); #else return getTime(); #endif diff --git a/c10/util/generic_math.h b/c10/util/generic_math.h index 9c2a424560484..e98be901c8c53 100644 --- a/c10/util/generic_math.h +++ b/c10/util/generic_math.h @@ -59,6 +59,10 @@ inline C10_HOST_DEVICE scalar_t div_floor_floating(scalar_t a, scalar_t b) template inline C10_HOST_DEVICE scalar_t div_floor_integer(scalar_t a, scalar_t b) { + if (C10_UNLIKELY(b == 0)) { + return scalar_t(0); + } + if (C10_UNLIKELY( std::is_signed::value && a == std::numeric_limits::min() && b == scalar_t(-1))) { diff --git a/c10/xpu/XPUCachingAllocator.cpp b/c10/xpu/XPUCachingAllocator.cpp index 7e534c7aa403b..f01e0c5c15224 100644 --- a/c10/xpu/XPUCachingAllocator.cpp +++ b/c10/xpu/XPUCachingAllocator.cpp @@ -1547,7 +1547,7 @@ class DeviceCachingAllocator { std::shared_ptr context) { if (!record_history) return; - bool should_skip = skip_actions_list.count(action) > 0; + bool should_skip = skip_actions_list.contains(action); if (should_skip) return; TraceEntry te( diff --git a/c10/xpu/XPUStream.h b/c10/xpu/XPUStream.h index c922759c2c48b..396002eb211b9 100644 --- a/c10/xpu/XPUStream.h +++ b/c10/xpu/XPUStream.h @@ -109,6 +109,11 @@ class C10_XPU_API XPUStream { } } + bool is_capturing() const { + return queue().ext_oneapi_get_state() == + sycl::ext::oneapi::experimental::queue_state::recording; + } + /// Return the priority that this stream is associated with. Lower numbers /// represent higher priority. int priority() const; diff --git a/c10/xpu/impl/XPUGuardImpl.h b/c10/xpu/impl/XPUGuardImpl.h index ad788ceb88141..d9a3d83ee4488 100644 --- a/c10/xpu/impl/XPUGuardImpl.h +++ b/c10/xpu/impl/XPUGuardImpl.h @@ -50,7 +50,11 @@ struct XPUGuardImpl final : public c10::impl::DeviceGuardImplInterface { cap.capability_data.capability_bits = (1ULL << kIndex_Byte) | (1ULL << kIndex_Char) | (1ULL << kIndex_Short) | (1ULL << kIndex_Int) | (1ULL << kIndex_Long) | (1ULL << kIndex_Float) | - (1ULL << kIndex_ComplexFloat) | (1ULL << kIndex_Bool); + (1ULL << kIndex_ComplexFloat) | (1ULL << kIndex_Bool) | + (1ULL << kIndex_Float8_e5m2) | (1ULL << kIndex_Float8_e4m3fn) | + (1ULL << kIndex_Float8_e5m2fnuz) | (1ULL << kIndex_Float8_e4m3fnuz) | + (1ULL << kIndex_Float8_e8m0fnu) | (1ULL << kIndex_UInt16) | + (1ULL << kIndex_UInt32) | (1ULL << kIndex_UInt64); // BFloat16 may be emulated. We always assume BFloat16 is available; // users can call is_bf16_supported() to check for native hardware support. cap.capability_data.capability_bits |= (1ULL << kIndex_BFloat16); @@ -214,6 +218,11 @@ struct XPUGuardImpl final : public c10::impl::DeviceGuardImplInterface { xpu_stream.synchronize(); } + bool isStreamCapturing(const Stream& stream) const override { + const XPUStream xpu_stream{stream}; + return xpu_stream.is_capturing(); + } + void synchronizeEvent(void* event) const override { if (!event) return; diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt index 6d6edbfa9e5ae..171e5446f5ed6 100644 --- a/caffe2/CMakeLists.txt +++ b/caffe2/CMakeLists.txt @@ -53,6 +53,9 @@ if(INTERN_BUILD_ATEN_OPS) # Add source, includes, and libs to lists list(APPEND Caffe2_CPU_SRCS ${ATen_CPU_SRCS}) + if(USE_MTIA) + list(APPEND Caffe2_CPU_SRCS ${ATen_MTIA_SRCS}) + endif() list(APPEND Caffe2_GPU_SRCS ${ATen_CUDA_CPP_SRCS}) list(APPEND Caffe2_XPU_SRCS ${ATen_XPU_SRCS}) list(APPEND Caffe2_XPU_INCLUDE ${ATen_XPU_INCLUDE}) @@ -556,13 +559,6 @@ if(USE_CUDA OR USE_ROCM) append_filelist("libtorch_cuda_core_sources" Caffe2_GPU_HIP_JIT_FUSERS_SRCS) endif() -# NativeRT is disabled -# if(USE_CUDA) -# append_filelist("libtorch_nativert_cuda_sources" Caffe2_GPU_SRCS) -# endif() -# if(USE_ROCM) -# append_filelist("libtorch_nativert_cuda_sources" Caffe2_HIP_SRCS) -# endif() if(USE_CUDA) list(APPEND Caffe2_GPU_CU_SRCS ${Caffe2_GPU_HIP_JIT_FUSERS_SRCS}) @@ -594,6 +590,7 @@ if(USE_CUDA) ${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp ${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.cu ${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/nccl_extension.cu + ${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/ops/nccl_reduce_scatter_offset.cu ${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/cuda_mem_pool.cpp PROPERTIES COMPILE_FLAGS "-DPYTORCH_C10_DRIVER_API_SUPPORTED=1" ) @@ -604,7 +601,7 @@ if(USE_CUDA) if(CMAKE_COMPILER_IS_GNUCXX) set_source_files_properties(${ASYNC_MM_FILE} PROPERTIES COMPILE_FLAGS "-Wno-unused-but-set-variable") endif() - if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0 AND CUDA_NVCC_FLAGS MATCHES ".*compute_90.*") + if(CUDA_NVCC_FLAGS MATCHES ".*compute_90.*") set_source_files_properties(${ASYNC_MM_FILE} PROPERTIES COMPILE_FLAGS "-gencode arch=compute_90a,code=sm_90a") endif() endif() @@ -845,6 +842,10 @@ if(USE_MPS) if(CAN_COMPILE_METAL) add_dependencies(torch_cpu metallibs) target_link_options(torch_cpu PRIVATE -Wl,-sectcreate,__TEXT,metal_basic,${CMAKE_CURRENT_BINARY_DIR}/aten/src/ATen/kernels_basic.metallib) + if(CAN_COMPILE_METAL_40) + target_link_options(torch_cpu PRIVATE -Wl,-sectcreate,__TEXT,metal_40,${CMAKE_CURRENT_BINARY_DIR}/aten/src/ATen/kernels_40.metallib) + target_compile_definitions(torch_cpu PRIVATE CAN_BUILD_METAL_4) + endif() else() target_compile_definitions(torch_cpu PRIVATE PYTORCH_JIT_COMPILE_SHADERS) endif() @@ -1369,8 +1370,6 @@ if(BUILD_TEST) else() add_subdirectory(${TORCH_ROOT}/test/cpp/jit ${CMAKE_BINARY_DIR}/test_jit) add_subdirectory(${TORCH_ROOT}/test/cpp/lazy ${CMAKE_BINARY_DIR}/test_lazy) - # NativeRT is disabled - # add_subdirectory(${TORCH_ROOT}/test/cpp/nativert ${CMAKE_BINARY_DIR}/test_nativert) add_subdirectory(${TORCH_ROOT}/test/cpp/profiler ${CMAKE_BINARY_DIR}/test_profiler) add_subdirectory(${TORCH_ROOT}/test/inductor ${CMAKE_BINARY_DIR}/test_inductor) add_subdirectory(${TORCH_ROOT}/test/cpp/aoti_abi_check ${CMAKE_BINARY_DIR}/test_aoti_abi_check) @@ -1468,6 +1467,74 @@ if(USE_ROCM) if(USE_ROCM_CK_SDPA) target_compile_definitions(torch_hip PRIVATE USE_ROCM_CK_SDPA) endif() + + message(INFO "USE_NVSHMEM=${USE_NVSHMEM} USE_ROCM=${USE_ROCM} rocshmem_FOUND=${rocshmem_FOUND}") + if(USE_NVSHMEM) + # rocSHMEM upstream ships device code for a subset of AMD gfx targets only. + set(ROCSHMEM_SUPPORTED_ARCH gfx90a gfx942 gfx950 gfx1100 gfx1201) + set(_torch_rocshmem_build_arches) + foreach(_arch ${PYTORCH_ROCM_ARCH}) + list(FIND ROCSHMEM_SUPPORTED_ARCH "${_arch}" _rocsmem_supported_idx) + if(NOT _rocsmem_supported_idx EQUAL -1) + # rocSHMEM device archives are commonly built for gfx90a xnack variants, + # so map plain gfx90a to gfx90a:xnack-/gfx90a:xnack+ to match device + # symbols during HIP RDC link. + if(_arch STREQUAL "gfx90a") + list(APPEND _torch_rocshmem_build_arches "gfx90a:xnack-" "gfx90a:xnack+") + else() + list(APPEND _torch_rocshmem_build_arches "${_arch}") + endif() + endif() + endforeach() + if(rocshmem_FOUND AND _torch_rocshmem_build_arches) + message(STATUS "rocSHMEM found, building with rocSHMEM support for arches: ${_torch_rocshmem_build_arches}") + message(STATUS "ROCSHMEM_INCLUDE_DIR: '${ROCSHMEM_INCLUDE_DIR}'") + set(TORCH_ROCSHMEM_SRCS + "${TORCH_SRC_DIR}/csrc/distributed/c10d/cuda/utils.cpp" + "${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/rocshmem_extension.cu" + "${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/NVSHMEMSymmetricMemory.cpp" + "${TORCH_SRC_DIR}/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.cpp" + ) + set_source_files_properties(${TORCH_ROCSHMEM_SRCS} PROPERTIES LANGUAGE HIP) + + # FindHIP's generated compile script bakes in global HIP_CLANG_FLAGS. Override + # -fno-gpu-rdc with -fgpu-rdc so device linker can pull in rocSHMEM device APIs. + set(HIP_CLANG_FLAGS_SAVED ${HIP_CLANG_FLAGS}) + set(HIP_CLANG_FLAGS_FOR_ROCSHMEM ${HIP_CLANG_FLAGS}) + list(FILTER HIP_CLANG_FLAGS_FOR_ROCSHMEM EXCLUDE REGEX "^-fno-gpu-rdc$") + list(APPEND HIP_CLANG_FLAGS_FOR_ROCSHMEM -fgpu-rdc) + set(HIP_CLANG_FLAGS ${HIP_CLANG_FLAGS_FOR_ROCSHMEM}) + + hip_add_library(torch_rocshmem SHARED ${TORCH_ROCSHMEM_SRCS}) + + set(HIP_CLANG_FLAGS ${HIP_CLANG_FLAGS_SAVED}) + + set_target_properties(torch_rocshmem PROPERTIES LINKER_LANGUAGE HIP) + + + # torch_rocshmem is built with -fgpu-rdc, so its HIP link step performs + # device linking and must carry the full arch list explicitly. + set_target_properties( + torch_rocshmem PROPERTIES HIP_ARCHITECTURES "${_torch_rocshmem_build_arches}") + + + target_compile_definitions(torch_hip PUBLIC USE_NVSHMEM) + target_compile_definitions(torch_rocshmem PUBLIC USE_NVSHMEM USE_ROCM) + + target_link_libraries(torch_rocshmem PRIVATE roc::rocshmem) + target_link_libraries(torch_hip PRIVATE torch_rocshmem) + + install(TARGETS torch_rocshmem EXPORT Caffe2Targets DESTINATION + "${TORCH_INSTALL_LIB_DIR}") + elseif(rocshmem_FOUND) + message(STATUS + "rocSHMEM found but skipped: PYTORCH_ROCM_ARCH has no entry in " + "ROCSHMEM_SUPPORTED_ARCH (${ROCSHMEM_SUPPORTED_ARCH}). " + "Building without rocSHMEM support.") + else() + message(WARNING "rocSHMEM not found (USE_NVSHMEM=ON). Building without rocSHMEM support.") + endif() #rocshmem_FOUND / arch intersection + endif() #USE_NVSHMEM endif() if(BUILD_LITE_INTERPRETER) @@ -1990,7 +2057,11 @@ if(BUILD_TEST) endif() if(USE_ROCM) + string(JOIN " " HIP_HIPCC_FLAGS_STR ${HIP_HIPCC_FLAGS}) + set(HIP_HIPCC_FLAGS ${HIP_HIPCC_FLAGS_STR}) + set(BASE_HIPCC_FLAGS ${HIP_HIPCC_FLAGS}) + foreach(test_src ${Caffe2_HIP_TEST_SRCS}) get_filename_component(test_name ${test_src} NAME_WE) if(test_src MATCHES "^.*\.hip$") diff --git a/caffe2/perfkernels/batch_box_cox_avx512.cc b/caffe2/perfkernels/batch_box_cox_avx512.cc deleted file mode 100644 index a97cb364a359d..0000000000000 --- a/caffe2/perfkernels/batch_box_cox_avx512.cc +++ /dev/null @@ -1,118 +0,0 @@ -#ifdef CAFFE2_PERF_USE_MKL -#include - -// Enable compiler vectorized version only if numerical consistency is not -// required between dev and opt versions - disabled for now -#ifndef FAST_VECTORIZED_KERNEL -#define CPU_CAPABILITY_AVX512 -#include - -namespace at::vec { -namespace { -// Implements the vectorized version of std::max() operation, -// which DOESNOT propagates NaN for second argument -template -Vectorized max(const Vectorized& a, const Vectorized& b); - -template <> -Vectorized max(const Vectorized& a, const Vectorized& b) { - // std::max(NaN, nonNan) -> NaN - return _mm512_max_pd(b, a); -} - -template <> -Vectorized max(const Vectorized& a, const Vectorized& b) { - // std::max(NaN, nonNan) -> NaN - return _mm512_max_ps(b, a); -} - -// Implements recieprocal method based on newton-rapson method -// 1. user RCP approximiation -// 2. update with RCP = RCP * (2 - X * RCP) -template -Vectorized fast_recieprocal(const Vectorized& b); -template -scalar_t fast_recieprocal(scalar_t b); - -template<> -Vectorized fast_recieprocal(const Vectorized& b) { - auto minus2 = _mm512_set1_ps(-2.f); - auto rcp = _mm512_rcp14_ps(b); - rcp = _mm512_mul_ps(rcp, _mm512_fnmsub_ps(rcp, b, minus2)); - rcp = _mm512_mul_ps(rcp, _mm512_fnmsub_ps(rcp, b, minus2)); - return rcp; -} - -template <> -float fast_recieprocal(float b) { - auto minus2 = _mm_set_ss(-2.f); - auto b_reg = _mm_set_ss(b); - auto rcp = _mm_rcp_ss(b_reg); - rcp = _mm_mul_ss(rcp, _mm_fnmsub_ss(rcp, b_reg, minus2)); - rcp = _mm_mul_ss(rcp, _mm_fnmsub_ss(rcp, b_reg, minus2)); - return _mm_cvtss_f32(rcp); -} - -template<> -Vectorized fast_recieprocal(const Vectorized& b) { - auto minus2 = _mm512_set1_pd(-2.); - auto rcp = _mm512_rcp14_pd(b); - rcp = _mm512_mul_pd(rcp, _mm512_fnmsub_pd(rcp, b, minus2)); - rcp = _mm512_mul_pd(rcp, _mm512_fnmsub_pd(rcp, b, minus2)); - return rcp; -} - -template <> -double fast_recieprocal(double b) { - return 1./b; -} -} // namespace -} // namespace at::vec -#endif - -#include "caffe2/perfkernels/batch_box_cox_vec.h" - -namespace caffe2::details { - -template -void compute_batch_box_cox__avx512( - std::size_t N, - std::size_t D, - std::size_t block_size, - const T* self_data, - const T* __restrict lambda1_data, - const T* __restrict lambda2_data, - T* output_data) { - compute_batch_box_cox_vec_fma( - N, - D, - block_size, - self_data, - lambda1_data, - lambda2_data, - output_data); - } - -// Vectorized version specializations for float and double -template -void compute_batch_box_cox__avx512( - std::size_t N, - std::size_t D, - std::size_t block_size, - const float* self_data, - const float* __restrict lambda1_data, - const float* __restrict lambda2_data, - float* output_data); - -template -void compute_batch_box_cox__avx512( - std::size_t N, - std::size_t D, - std::size_t block_size, - const double* self_data, - const double* __restrict lambda1_data, - const double* __restrict lambda2_data, - double* output_data); - -} // namespace caffe2::detail -#endif // CAFFE2_PERF_USE_MKL diff --git a/caffe2/perfkernels/batch_box_cox_sve128.cc b/caffe2/perfkernels/batch_box_cox_sve128.cc deleted file mode 100644 index 897e3a8ee4755..0000000000000 --- a/caffe2/perfkernels/batch_box_cox_sve128.cc +++ /dev/null @@ -1,200 +0,0 @@ -#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) && defined(CAFFE2_PERF_WITH_SVE128) -#include -#include -#include -#include -#include - -#include "c10/macros/Macros.h" - -/// Select `svlog` accuracy: -/// - 0: original. -/// - 1: more accurate, similar performance. -/// - 2: very high accuracy, a bit lower speed. -#define SVLOG_ACCURACY 2 - -/// Handle special cases in `svexp`: -/// - 0: original. -/// - 1: use clamp, better performance. -/// - 2: no special case handling. -#define SVEXP_SPECIAL_CLAMP 1 - -#if SVLOG_ACCURACY == 2 -static inline svfloat32_t svlog(svfloat32_t x) { - const svbool_t ptrue = svptrue_b8(); - - svint32_t u = svreinterpret_s32(x) - 0x3F2AAAAB; - - svfloat32_t r = svreinterpret_f32((u & 0x007FFFFF) + 0x3F2AAAAB) - 1.0f; - svfloat32_t n = svcvt_f32_x(ptrue, u >> 23); - asm("" : "+w"(r)); // NOTE: can improve instruction scheduling. - - svfloat32_t r2 = r * r; - svfloat32_t p = -0x1.4F9934p-3f + r * 0x1.5A9AA2p-3f; - svfloat32_t q = -0x1.00187Cp-2f + r * 0x1.961348p-3f; - svfloat32_t y = -0x1.FFFFC8p-2f + r * 0x1.555D7Cp-2f; - return (r + n * 0x1.62E43p-1f) + - (y + (q + (p + -0x1.3E737Cp-3f * r2) * r2) * r2) * r2; -} -#elif SVLOG_ACCURACY == 1 -static inline svfloat32_t svlog(svfloat32_t x) { - const svbool_t ptrue = svptrue_b8(); - - svint32_t u = svreinterpret_s32(x) - 0x3F2AAAAB; - - svfloat32_t r = svreinterpret_f32((u & 0x007FFFFF) + 0x3F2AAAAB) - 1.0f; - svfloat32_t n = svcvt_f32_x(ptrue, u >> 23); - asm("" : "+w"(r)); // NOTE: can improve instruction scheduling. - - svfloat32_t r2 = r * r; - svfloat32_t A = -0x1.923814p-3f + r * 0x1.689E5Ep-3f; - svfloat32_t B = -0x1.FC0968p-3f + r * 0x1.93BF0Cp-3f; - svfloat32_t C = -0x1.000478p-1f + r * 0x1.556906p-2f; - - return (r + n * 0x1.62E43p-1f) + (C + (B + A * r2) * r2) * r2; -} -#elif SVLOG_ACCURACY == 0 -static inline svfloat32_t svlog(svfloat32_t x) { - const svbool_t ptrue = svptrue_b8(); - - svint32_t u = svsra_n_s32(svdup_n_s32(-127), svreinterpret_s32(x), 23); - - svfloat32_t n = svcvt_f32_x(ptrue, u); - svfloat32_t r = svreinterpret_f32(svreinterpret_s32(x) - (u << 23)); - - svfloat32_t D = -0.165253549814f + r * 0.0141278216615f; - svfloat32_t C = -2.47071170807f + r * 0.844007015228f; - svfloat32_t B = -5.68692588806f + r * 4.58445882797f; - svfloat32_t A = -2.29561495781f + r * 5.17591238022f; - - svfloat32_t r2 = r * r; - return (A + n * 0.6931471805f) + (B + (C + D * r2) * r2) * r2; -} -#endif - -static inline svfloat32_t svexp(svfloat32_t x) { - // Clamp interval set to prevent denormals! - const svfloat32_t max_input = svdup_n_f32(88.722839f); - const svfloat32_t min_input = svdup_n_f32(-87.33654f); - const svfloat32_t shift = svdup_n_f32(0x1.0000FEp+23f); - const svbool_t ptrue = svptrue_b8(); - -#if SVEXP_SPECIAL_CLAMP == 1 - x = svmax_x(ptrue, svmin_x(ptrue, x, max_input), min_input); -#endif - - svfloat32_t z = svmla_n_f32_x(ptrue, shift, x, 0x1.715476p+0f); - svfloat32_t n = z - shift; - svfloat32_t scale = svreinterpret_f32(svreinterpret_u32(z) << 23); - - svfloat32_t r_hi = x - n * 0x1.62E400p-1f; - svfloat32_t r = r_hi - n * 0x1.7F7D1Cp-20f; - svfloat32_t r2 = r * r; - - svfloat32_t C = 0x1.573E2Ep-5f + r * 0x1.0E4020p-7f; - svfloat32_t B = 0x1.FFFDB6p-2f + r * 0x1.555E66p-3f; - svfloat32_t A = r * 0x1.FFFFECp-1f; - - svfloat32_t poly = scale + (A + (B + C * r2) * r2) * scale; - -#if SVEXP_SPECIAL_CLAMP == 0 - const svfloat32_t inf = svdup_n_f32(std::numeric_limits::infinity()); - poly = svsel_f32(svcmplt_f32(ptrue, x, min_input), svdup_n_f32(0.0f), poly); - poly = svsel_f32(svcmpgt_f32(ptrue, x, max_input), inf, poly); -#endif - - return poly; -} - -static inline svfloat32_t compute_batch_box_cox_vec_sve128_float( - svfloat32_t lambda1_v, - svfloat32_t lambda2_v, - svfloat32_t data_v, - svfloat32_t k_eps) { - const svbool_t ptrue = svptrue_b8(); - - svfloat32_t lnData = svlog(svmax_x(ptrue, data_v + lambda2_v, k_eps)); - svbool_t predNZ = svcmpne_n_f32(ptrue, lambda1_v, 0.0f); - if (C10_LIKELY(svptest_any(predNZ, predNZ))) { - svfloat32_t lambda1_r = svdivr_f32_m(predNZ, lambda1_v, svdup_n_f32(1.0f)); - svfloat32_t pow = svexp(lnData * lambda1_v); - lnData = svsel_f32(predNZ, lambda1_r, lnData); - lnData = svnmsb_f32_m(predNZ, lnData, pow, lnData); - } - return lnData; -} - -template -void compute_batch_box_cox_vec_sve128( - std::size_t N, - std::size_t D, - const T* data_ptr, - const T* __restrict lambda1_ptr, - const T* __restrict lambda2_ptr, - T* output_ptr); - -template <> -void compute_batch_box_cox_vec_sve128( - std::size_t N, - std::size_t D, - const float *data_ptr, - const float *__restrict lambda1_ptr, - const float *__restrict lambda2_ptr, - float *output_ptr) { - const svfloat32_t k_eps = svdup_n_f32(static_cast(1e-6)); - - std::size_t remainder = D % 4; - std::size_t loopBound = D - remainder; - svbool_t remainderPred = svwhilelt_b32_u64(0, remainder); - - for (; C10_LIKELY(N > 0); --N) { - for (std::size_t j = 0; C10_LIKELY(j != loopBound); - j += 4, data_ptr += 4, output_ptr += 4) { - svfloat32_t lambda1_v = - svset_neonq(svundef_f32(), vld1q_f32(lambda1_ptr + j)); - svfloat32_t lambda2_v = - svset_neonq(svundef_f32(), vld1q_f32(lambda2_ptr + j)); - svfloat32_t data_v = svset_neonq(svundef_f32(), vld1q_f32(data_ptr)); - svfloat32_t result = compute_batch_box_cox_vec_sve128_float( - lambda1_v, lambda2_v, data_v, k_eps); - vst1q_f32(output_ptr, svget_neonq(result)); - } - if (C10_LIKELY(remainder > 0)) { - svfloat32_t lambda1_v = svld1_f32(remainderPred, lambda1_ptr + loopBound); - svfloat32_t lambda2_v = svld1_f32(remainderPred, lambda2_ptr + loopBound); - svfloat32_t data_v = svld1_f32(remainderPred, data_ptr); - svfloat32_t result = compute_batch_box_cox_vec_sve128_float( - lambda1_v, lambda2_v, data_v, k_eps); - svst1_f32(remainderPred, output_ptr, result); - data_ptr += remainder; - output_ptr += remainder; - } - } -} - -namespace caffe2::details { - -template -void compute_batch_box_cox__sve128( - std::size_t N, - std::size_t D, - const T* self_data, - const T* __restrict lambda1_data, - const T* __restrict lambda2_data, - T* output_data) { - compute_batch_box_cox_vec_sve128( - N, D, self_data, lambda1_data, lambda2_data, output_data); -} - -// Vectorized version specializations for float and double -template void compute_batch_box_cox__sve128( - std::size_t N, - std::size_t D, - const float* self_data, - const float* __restrict lambda1_data, - const float* __restrict lambda2_data, - float* output_data); - -} // namespace caffe2::details - -#endif // __aarch64__ && __ARM_FEATURE_SVE && CAFFE2_PERF_WITH_SVE128 diff --git a/caffe2/perfkernels/batch_box_cox_vec.h b/caffe2/perfkernels/batch_box_cox_vec.h deleted file mode 100644 index 08e4f84fe4327..0000000000000 --- a/caffe2/perfkernels/batch_box_cox_vec.h +++ /dev/null @@ -1,321 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include "vectorizer.h" -#include - -namespace caffe2::details { - -namespace { -void TileIndicesInPlace(std::vector& v, const std::size_t D, const std::size_t K) { - auto n = v.size(); - v.resize(K * n); - for (const auto k : c10::irange(1, K)) { - for (const auto j : c10::irange(n)) { - v[k * n + j] = v[j] + k * D; - } - } -} - -// MKL VML function templates. -template -void PackV(const int N, const T* a, const int* ia, T* y); -template -void UnpackV(const int N, const T* a, T* y, const int* iy); - -#define DELEGATE_PACKV_FUNCTION(T, OriginalFunc) \ - template <> \ - void PackV(const int N, const T* a, const int* ia, T* y) { \ - OriginalFunc(N, a, ia, y); \ - } -DELEGATE_PACKV_FUNCTION(float, vsPackV) -DELEGATE_PACKV_FUNCTION(double, vdPackV) -#undef DELEGATE_PACKV_FUNCTION - -#define DELEGATE_UNPACKV_FUNCTION(T, OriginalFunc) \ - template <> \ - void UnpackV(const int N, const T* a, T* y, const int* iy) { \ - OriginalFunc(N, a, y, iy); \ - } -DELEGATE_UNPACKV_FUNCTION(float, vsUnpackV) -DELEGATE_UNPACKV_FUNCTION(double, vdUnpackV) -#undef DELEGATE_UNPACKV_FUNCTION - -#ifndef FAST_VECTORIZED_KERNEL -template -void box_cox_zero_lambda( - size_t D, - const T* const self_data, - const T* const lambda2_data, - T k_eps, - T* const output_data) { - int j = 0; - using Vec = at::vec::Vectorized; - constexpr int64_t VLEN = Vec::size(); - auto k_eps_vec = Vec(k_eps); - for(; j + VLEN < D; j += VLEN) { - auto data = Vec::loadu(self_data + j); - auto lambda2 = Vec::loadu(lambda2_data + j); - auto sum = data + lambda2; - auto max = at::vec::max(sum, k_eps_vec); - auto res = max.log(); - res.store(output_data + j); - } - for ( ;j < D; ++j) { - auto sum = self_data[j] + lambda2_data[j]; - auto max = std::max(sum, k_eps); - output_data[j] = std::log(max); - } -} - -template -at::vec::Vectorized box_cox_nonzero_lambda_impl( - at::vec::Vectorized data, - at::vec::Vectorized lambda1, - at::vec::Vectorized lambda2, - at::vec::Vectorized k_eps) { - auto sum = data + lambda2; - auto max = at::vec::max(sum, k_eps); - auto lambda_over_1 = at::vec::fast_recieprocal(lambda1); - auto pow = max.pow(lambda1); - return at::vec::fmsub(pow, lambda_over_1, lambda_over_1); -} - -template -void box_cox_nonzero_lambda( - int64_t D, - const T* data_ptr, - const T* lambda1_ptr, - const T* lambda2_ptr, - T k_eps, - T* out) { - - int j = 0; - using Vec = at::vec::Vectorized; - constexpr int64_t VLEN = Vec::size(); - auto k_eps_vec = Vec(k_eps); - for(; j + VLEN < D; j += VLEN) { - auto data = Vec::loadu(data_ptr + j); - auto lambda1 = Vec::loadu(lambda1_ptr + j); - auto lambda2 = Vec::loadu(lambda2_ptr + j); - auto res = box_cox_nonzero_lambda_impl(data, lambda1, lambda2, k_eps_vec); - res.store(out + j); - } - if (j < D) { - auto remaining = D - j; - auto data = Vec::loadu(data_ptr + j, remaining); - auto lambda1 = Vec::loadu(lambda1_ptr + j, remaining); - auto lambda2 = Vec::loadu(lambda2_ptr + j, remaining); - auto res = box_cox_nonzero_lambda_impl(data, lambda1, lambda2, k_eps_vec); - res.store(out + j, remaining); - } -} -#else -template -void box_cox_zero_lambda( - size_t D, - const T* const self_data, - const T* const lambda2_data, - T k_eps, - T* const output_data) { - VECTOR_LOOP for (auto j=0 ;j < D; ++j) { - auto sum = self_data[j] + lambda2_data[j]; - auto max = std::max(sum, k_eps); - output_data[j] = std::log(max); - } -} - -template -void box_cox_nonzero_lambda( - int64_t D, - const T* data_ptr, - const T* lambda1_ptr, - const T* lambda2_ptr, - T k_eps, - T* out) { - - VECTOR_LOOP for (auto j=0 ;j < D; ++j) { - FAST_MATH - auto sum = data_ptr[j] + lambda2_ptr[j]; - auto max = std::max(sum, k_eps); - auto lamda1 = lambda1_ptr[j]; - auto lambda_over_1 = 1 / lamda1; - if constexpr (std::is_same::value) { - lambda_over_1 = lambda_over_1 * (T{2} - lambda_over_1 * lamda1); - lambda_over_1 = lambda_over_1 * (T{2} - lambda_over_1 * lamda1); - } - auto pow = std::pow(max, lamda1); - out[j] = pow * lambda_over_1 - lambda_over_1; - } -} -#endif // FAST_VECTORIZED_KERNEL - -template -void box_cox_mixed_lambda( - const T* const self_data, - const std::vector& nonzeros, - const std::vector& zeros, - const T* const lambda1, - const T* const lambda2, - const T* const lambda2_z_, - T k_eps, - T* const buffer, - T* const output_data) { - PackV(nonzeros.size(), self_data, nonzeros.data(), buffer); - box_cox_nonzero_lambda( - nonzeros.size(), buffer, lambda1, lambda2, k_eps, buffer); - UnpackV(nonzeros.size(), buffer, output_data, nonzeros.data()); - - PackV(zeros.size(), self_data, zeros.data(), buffer); - box_cox_zero_lambda( - zeros.size(), buffer, lambda2_z_, k_eps, buffer); - UnpackV(zeros.size(), buffer, output_data, zeros.data()); -} - -template -void TileArrayIntoVector( - const T* const a, - const size_t D, - const int K, - std::vector& b) { - b.resize(K * D); - for (const auto k : c10::irange(K)) { - std::copy(a, a + D, b.begin() + k * D); - } -} - -template -void compute_batch_box_cox_vec_fma( - std::size_t N, - std::size_t D, - std::size_t block_size, - const T* self_data, - const T* __restrict lambda1_data, - const T* __restrict lambda2_data, - T* output_data) { - constexpr T k_eps = static_cast(1e-6); - - FOLLY_DECLARE_REUSED(zeros, std::vector); - FOLLY_DECLARE_REUSED(nonzeros, std::vector); - // Don't bother calling reserve; calls after the first will get a - // correctly-sized allocation anyway. - for (const auto j : c10::irange(D)) { - if (lambda1_data[j] == 0) { - zeros.push_back(j); - } else { - nonzeros.push_back(j); - } - } - - // Process K rows at a time for effective vectorization with small rows. - const auto K = std::min(N, (block_size + D - 1) / D); - - FOLLY_DECLARE_REUSED(lambda1_, std::vector); - FOLLY_DECLARE_REUSED(lambda2_, std::vector); - FOLLY_DECLARE_REUSED(lambda2_z_, std::vector); - - if (nonzeros.size() == D) { - // ((x + lambda2)^lambda1 - 1)/lambda1, if lambda1 != 0 - size_t i = 0; - if (K > 1) { - TileArrayIntoVector(lambda1_data, D, K, lambda1_); - TileArrayIntoVector(lambda2_data, D, K, lambda2_); - DCHECK_EQ(K * D, lambda1_.size()); - DCHECK_EQ(K * D, lambda2_.size()); - for (; i < N - K + 1; i += K, self_data += K * D, output_data += K * D) { - box_cox_nonzero_lambda( - K * D, - self_data, - lambda1_.data(), - lambda2_.data(), - k_eps, - output_data); - } - } - for (; i < N; i++, self_data += D, output_data += D) { - box_cox_nonzero_lambda( - D, self_data, lambda1_data, lambda2_data, k_eps, output_data); - } - } else if (zeros.size() == D) { - // ln(x + lambda2), if lambda1 == 0 - size_t i = 0; - if (K > 1) { - TileArrayIntoVector(lambda2_data, D, K, lambda2_z_); - DCHECK_EQ(K * D, lambda2_z_.size()); - for (; i < N - K + 1; i += K, self_data += K * D, output_data += K * D) { - box_cox_zero_lambda( - K * D, self_data, lambda2_z_.data(), k_eps, output_data); - } - } - for (; i < N; i++, self_data += D, output_data += D) { - box_cox_zero_lambda( - D, self_data, lambda2_data, k_eps, output_data); - } - } else { - // mix zeros and nonzeros - const size_t n = nonzeros.size(); - if (K > 1) { - TileIndicesInPlace(nonzeros, 0, K); - TileIndicesInPlace(zeros, 0, K); - } - - FOLLY_DECLARE_REUSED(buffer, std::vector); - - buffer.resize(std::max(nonzeros.size(), zeros.size())); - lambda1_.resize(nonzeros.size()); - lambda2_.resize(nonzeros.size()); - lambda2_z_.resize(zeros.size()); - PackV(nonzeros.size(), lambda1_data, nonzeros.data(), lambda1_.data()); - PackV(nonzeros.size(), lambda2_data, nonzeros.data(), lambda2_.data()); - PackV(zeros.size(), lambda2_data, zeros.data(), lambda2_z_.data()); - - size_t i = 0; - if (K > 1) { - // Truncate to original size, and re-tile with offsets this time. - nonzeros.resize(n); - DCHECK_GT(D, n); - zeros.resize(D - n); - TileIndicesInPlace(nonzeros, D, K); - TileIndicesInPlace(zeros, D, K); - DCHECK_EQ(nonzeros.size(), lambda1_.size()); - DCHECK_EQ(nonzeros.size(), lambda2_.size()); - DCHECK_EQ(zeros.size(), lambda2_z_.size()); - - for (; i < N - K + 1; i += K, self_data += K * D, output_data += K * D) { - box_cox_mixed_lambda( - self_data, - nonzeros, - zeros, - lambda1_.data(), - lambda2_.data(), - lambda2_z_.data(), - k_eps, - buffer.data(), - output_data); - } - // Truncate to original size. - nonzeros.resize(n); - zeros.resize(D - n); - } - for (; i < N; i++, self_data += D, output_data += D) { - box_cox_mixed_lambda( - self_data, - nonzeros, - zeros, - lambda1_.data(), - lambda2_.data(), - lambda2_z_.data(), - k_eps, - buffer.data(), - output_data); - } - } -} -} // namespace - -} // namespace caffe2::details diff --git a/cmake/Codegen.cmake b/cmake/Codegen.cmake index c3003e1c8af39..247337c32cf41 100644 --- a/cmake/Codegen.cmake +++ b/cmake/Codegen.cmake @@ -163,6 +163,11 @@ if(INTERN_BUILD_ATEN_OPS) set(GEN_XPU_FLAG --xpu) endif() + set(GEN_MTIA_FLAG) + if(USE_MTIA) + set(GEN_MTIA_FLAG --mtia) + endif() + set(CUSTOM_BUILD_FLAGS) if(INTERN_BUILD_MOBILE) if(USE_VULKAN) @@ -249,6 +254,7 @@ if(INTERN_BUILD_ATEN_OPS) ${GEN_ROCM_FLAG} ${GEN_MPS_FLAG} ${GEN_XPU_FLAG} + ${GEN_MTIA_FLAG} ${CUSTOM_BUILD_FLAGS} ) diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake index c39dd19c19252..70c5d3de4cf23 100644 --- a/cmake/Dependencies.cmake +++ b/cmake/Dependencies.cmake @@ -47,15 +47,21 @@ if(USE_CUDA) # torch::cudart is dealt with separately, due to CUDA_ADD_LIBRARY # design reason (it adds CUDA_LIBRARIES itself). set(Caffe2_PUBLIC_CUDA_DEPENDENCY_LIBS ) - if(NOT CAFFE2_USE_NVRTC) - caffe2_update_option(USE_NVRTC OFF) - endif() list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS caffe2::curand caffe2::cufft caffe2::cublas) if(CAFFE2_USE_CUDNN) + if(NOT CAFFE2_USE_NVRTC) + message(FATAL_ERROR + "USE_CUDNN requires USE_NVRTC (required by cudnn_frontend 1.21+). " + "Please set -DUSE_NVRTC=ON or disable cuDNN with -DUSE_CUDNN=OFF.") + endif() list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS torch::cudnn) + list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS caffe2::nvrtc_runtime) else() caffe2_update_option(USE_CUDNN OFF) endif() + if(NOT CAFFE2_USE_NVRTC) + caffe2_update_option(USE_NVRTC OFF) + endif() if(CAFFE2_USE_CUSPARSELT) list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS torch::cusparselt) else() @@ -1023,7 +1029,7 @@ if(USE_ROCM) list(APPEND HIP_CXX_FLAGS -DUSE_ROCM_CK_GEMM) endif() list(APPEND HIP_HIPCC_FLAGS --offload-compress) - list(APPEND HIP_HIPCC_FLAGS -std=c++17) + list(APPEND HIP_HIPCC_FLAGS -std=c++20) # Pass device library path for theRock nightly builds if(DEFINED ENV{HIP_DEVICE_LIB_PATH}) file(TO_CMAKE_PATH "$ENV{HIP_DEVICE_LIB_PATH}" _hip_device_lib_path) @@ -1063,6 +1069,8 @@ if(USE_ROCM) list(APPEND HIP_HIPCC_FLAGS -fclang-abi-compat=17) set(HIP_CLANG_FLAGS ${HIP_CXX_FLAGS}) + string(JOIN " " HIP_HIPCC_FLAGS_STR ${HIP_HIPCC_FLAGS}) + set(HIP_HIPCC_FLAGS ${HIP_HIPCC_FLAGS_STR}) set(CMAKE_HIP_FLAGS ${HIP_HIPCC_FLAGS}) # Ask hcc to generate device code during compilation so we can use # host linker to link. @@ -1092,6 +1100,13 @@ if(USE_ROCM) endif() endif() + # ROCM-SMI needed to support symmetric memory + if(USE_DISTRIBUTED AND UNIX) + list(APPEND Caffe2_PUBLIC_HIP_DEPENDENCY_LIBS + rocm_smi64 + ) + endif() + # ---[ Kernel asserts # Kernel asserts is disabled for ROCm by default. # It can be turned on by turning on the env USE_ROCM_KERNEL_ASSERT to the build system. diff --git a/cmake/EnvVarForwarding.cmake b/cmake/EnvVarForwarding.cmake new file mode 100644 index 0000000000000..a31ce6db26637 --- /dev/null +++ b/cmake/EnvVarForwarding.cmake @@ -0,0 +1,138 @@ +# Forward environment variables to CMake variables. +# +# This replicates the behavior of setup.py / tools/setup_helpers/cmake.py which +# passes all BUILD_*, USE_*, and CMAKE_* environment variables as -D flags, plus +# a set of additional variables that don't follow the prefix convention. + +# Additional env vars that are forwarded with a different CMake variable name. +set(_ENV_ALIASES + "CUDNN_LIB_DIR=CUDNN_LIBRARY" + "USE_CUDA_STATIC_LINK=CAFFE2_STATIC_LINK_CUDA" +) + +# Additional env vars forwarded with the same name. +set(_ENV_PASSTHROUGH + UBSAN_FLAGS + BLAS + WITH_BLAS + CUDA_HOST_COMPILER + CUDA_NVCC_EXECUTABLE + CUDA_SEPARABLE_COMPILATION + CUDNN_LIBRARY + CUDNN_INCLUDE_DIR + CUDNN_ROOT + EXPERIMENTAL_SINGLE_THREAD_POOL + INSTALL_TEST + INTEL_MKL_DIR + INTEL_OMP_DIR + MKL_THREADING + MKLDNN_CPU_RUNTIME + MSVC_Z7_OVERRIDE + CAFFE2_USE_MSVC_STATIC_RUNTIME + Numa_INCLUDE_DIR + Numa_LIBRARIES + ONNX_ML + ONNX_NAMESPACE + ATEN_THREADING + WERROR + OPENSSL_ROOT_DIR + STATIC_DISPATCH_BACKEND + SELECTED_OP_LIST + TORCH_CUDA_ARCH_LIST + TORCH_XPU_ARCH_LIST + TRACING_BASED + PYTHON_LIB_REL_PATH +) + +# Low-priority aliases: if the canonical var is not set, use the alias. +set(_LOW_PRIORITY_ALIASES + "CUDA_HOST_COMPILER=CMAKE_CUDA_HOST_COMPILER" + "CUDAHOSTCXX=CUDA_HOST_COMPILER" + "CMAKE_CUDA_HOST_COMPILER=CUDA_HOST_COMPILER" + "CMAKE_CUDA_COMPILER=CUDA_NVCC_EXECUTABLE" + "CUDACXX=CUDA_NVCC_EXECUTABLE" +) + +# Forward aliased env vars (env name -> different cmake name) +foreach(_alias IN LISTS _ENV_ALIASES) + string(REPLACE "=" ";" _parts "${_alias}") + list(GET _parts 0 _env_name) + list(GET _parts 1 _cmake_name) + if(DEFINED ENV{${_env_name}} AND NOT DEFINED ${_cmake_name}) + set(${_cmake_name} "$ENV{${_env_name}}" CACHE STRING "From env ${_env_name}" FORCE) + endif() +endforeach() + +# Forward passthrough env vars (same name) +foreach(_var IN LISTS _ENV_PASSTHROUGH) + if(DEFINED ENV{${_var}} AND NOT DEFINED ${_var}) + set(${_var} "$ENV{${_var}}" CACHE STRING "From env ${_var}" FORCE) + endif() +endforeach() + +# Forward all BUILD_*, USE_*, CMAKE_* env vars not already set as CMake +# variables, plus vars ending in EXITCODE or EXITCODE__TRYRUN_OUTPUT. +# This matches the existing behavior where setup.py passed everything with +# these prefixes/suffixes through to CMake. +# We use execute_process + env to get the full list since CMake has no +# built-in way to enumerate environment variables. +execute_process( + COMMAND "${CMAKE_COMMAND}" -E environment + OUTPUT_VARIABLE _all_env + OUTPUT_STRIP_TRAILING_WHITESPACE +) +string(REPLACE "\n" ";" _env_lines "${_all_env}") +foreach(_line IN LISTS _env_lines) + if(_line MATCHES "^([A-Za-z_0-9]+)=(.*)") + set(_var_name "${CMAKE_MATCH_1}") + set(_var_value "${CMAKE_MATCH_2}") + # Only forward vars with BUILD_/USE_/CMAKE_ prefix or *EXITCODE* suffix. + string(REGEX MATCH "^(BUILD_|USE_|CMAKE_)" _has_prefix "${_var_name}") + string(REGEX MATCH "(EXITCODE|EXITCODE__TRYRUN_OUTPUT)$" _has_suffix "${_var_name}") + if(NOT _has_prefix AND NOT _has_suffix) + continue() + endif() + if(NOT DEFINED ${_var_name}) + set(${_var_name} "${_var_value}" CACHE STRING "From environment" FORCE) + endif() + endif() +endforeach() + +# Low-priority aliases +foreach(_alias IN LISTS _LOW_PRIORITY_ALIASES) + string(REPLACE "=" ";" _parts "${_alias}") + list(GET _parts 0 _env_name) + list(GET _parts 1 _cmake_name) + if(DEFINED ENV{${_env_name}} AND NOT DEFINED ${_cmake_name}) + set(${_cmake_name} "$ENV{${_env_name}}" CACHE STRING "From env alias ${_env_name}" FORCE) + endif() +endforeach() + +# Ensure Python's purelib is on CMAKE_PREFIX_PATH so CMake can find +# packages installed there (e.g., pybind11, numpy). +if(Python_EXECUTABLE) + execute_process( + COMMAND "${Python_EXECUTABLE}" -c "import sysconfig; print(sysconfig.get_path('purelib'))" + OUTPUT_VARIABLE _py_purelib + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET + ) + if(_py_purelib AND NOT "${_py_purelib}" STREQUAL "") + list(PREPEND CMAKE_PREFIX_PATH "${_py_purelib}") + # Preserve paths from the CMAKE_PREFIX_PATH environment variable. + # Setting the cmake variable shadows the env var, so we must merge it in + # explicitly. This ensures conda's prefix (e.g. /opt/conda/envs/py_3.10) + # is present so cmake can find conda-provided libraries (libgomp, libnuma). + if(DEFINED ENV{CMAKE_PREFIX_PATH} AND NOT "$ENV{CMAKE_PREFIX_PATH}" STREQUAL "") + if(WIN32) + # On Windows the env var is already ;-separated and : appears in drive + # letters (e.g. C:\conda\envs\py310), so use it as-is. + set(_env_prefix "$ENV{CMAKE_PREFIX_PATH}") + else() + string(REPLACE ":" ";" _env_prefix "$ENV{CMAKE_PREFIX_PATH}") + endif() + list(APPEND CMAKE_PREFIX_PATH ${_env_prefix}) + endif() + list(REMOVE_DUPLICATES CMAKE_PREFIX_PATH) + endif() +endif() diff --git a/cmake/Metal.cmake b/cmake/Metal.cmake index c9565e2fc0e9e..f378dfb0e803d 100644 --- a/cmake/Metal.cmake +++ b/cmake/Metal.cmake @@ -45,6 +45,13 @@ set(BFLOAT_METAL_CODE " ptr[idx] += 1; } ") +set(LAMBDA_METAL_CODE " + kernel void test(device float* ptr, + uint idx [[thread_position_in_grid]]) { + auto fn = [](float x) { return x + 1.0; }; + ptr[idx] = fn(ptr[idx]); + } +") if(NOT CAN_COMPILE_METAL_FOUND) file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/bfloat_inc.metal" "${BFLOAT_METAL_CODE}") execute_process(COMMAND xcrun metal -std=metal3.1 bfloat_inc.metal @@ -59,6 +66,23 @@ if(NOT CAN_COMPILE_METAL_FOUND) message(WARNING "Machine can not compile metal shaders, fails with ${XCRUN_OUTPUT}") set(CAN_COMPILE_METAL NO CACHE BOOL "Host can compile metal shaders") endif() + if(CAN_COMPILE_METAL) + file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/lambda_test.metal" "${LAMBDA_METAL_CODE}") + execute_process(COMMAND xcrun metal -std=metal4.0 -c lambda_test.metal -o /dev/null + WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" + OUTPUT_VARIABLE XCRUN_OUTPUT + ERROR_VARIABLE XCRUN_OUTPUT + RESULT_VARIABLE XCRUN_RC) + if(${XCRUN_RC} EQUAL 0) + message(STATUS "Metal toolchain supports Metal 4.0") + set(CAN_COMPILE_METAL_40 YES CACHE BOOL "Host can compile Metal 4.0 shaders" FORCE) + else() + message(STATUS "Metal toolchain does not support Metal 4.0") + set(CAN_COMPILE_METAL_40 NO CACHE BOOL "Host can compile Metal 4.0 shaders" FORCE) + endif() + else() + set(CAN_COMPILE_METAL_40 NO CACHE BOOL "Host can compile Metal 4.0 shaders" FORCE) + endif() set(CAN_COMPILE_METAL_FOUND YES CACHE INTERNAL "Run check for shader compiler") endif() diff --git a/cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake b/cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake index cbc26c4b23ad8..642b96ba119e8 100644 --- a/cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake +++ b/cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake @@ -30,10 +30,10 @@ endif() set(CUDA_KNOWN_GPU_ARCHITECTURES "Kepler" "Maxwell") # This list will be used for CUDA_ARCH_NAME = Common option (enabled by default) -set(CUDA_COMMON_GPU_ARCHITECTURES "3.5" "5.0") +set(CUDA_COMMON_GPU_ARCHITECTURES "5.0") # This list is used to filter CUDA archs when autodetecting -set(CUDA_ALL_GPU_ARCHITECTURES "3.5" "5.0") +set(CUDA_ALL_GPU_ARCHITECTURES "5.0") if(CUDA_VERSION VERSION_GREATER "10.5") list(APPEND CUDA_KNOWN_GPU_ARCHITECTURES "Ampere") @@ -65,19 +65,10 @@ if(NOT CUDA_VERSION VERSION_LESS "11.8") list(APPEND CUDA_ALL_GPU_ARCHITECTURES "8.9") list(APPEND CUDA_ALL_GPU_ARCHITECTURES "9.0") - if(CUDA_VERSION VERSION_LESS "12.0") - set(CUDA_LIMIT_GPU_ARCHITECTURE "9.0") - list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "8.9+PTX") - list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "9.0+PTX") - endif() endif() -if(NOT CUDA_VERSION VERSION_LESS "12.0") - list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "9.0a") - list(APPEND CUDA_ALL_GPU_ARCHITECTURES "9.0a") - list(REMOVE_ITEM CUDA_COMMON_GPU_ARCHITECTURES "3.5") - list(REMOVE_ITEM CUDA_ALL_GPU_ARCHITECTURES "3.5") -endif() +list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "9.0a") +list(APPEND CUDA_ALL_GPU_ARCHITECTURES "9.0a") if(CUDA_VERSION VERSION_GREATER "12.6") list(APPEND CUDA_KNOWN_GPU_ARCHITECTURES "Blackwell") diff --git a/cmake/Summary.cmake b/cmake/Summary.cmake index 686f5861960d3..2d68d81c6b367 100644 --- a/cmake/Summary.cmake +++ b/cmake/Summary.cmake @@ -149,6 +149,7 @@ function(caffe2_print_configuration_summary) message(STATUS " USE_PYTORCH_METAL_EXPORT : ${USE_PYTORCH_METAL_EXPORT}") message(STATUS " USE_MPS : ${USE_MPS}") message(STATUS " CAN_COMPILE_METAL : ${CAN_COMPILE_METAL}") + message(STATUS " CAN_COMPILE_METAL_40 : ${CAN_COMPILE_METAL_40}") message(STATUS " USE_MKL : ${CAFFE2_USE_MKL}") if(${CAFFE2_USE_MKL}) message(STATUS " USE_STATIC_MKL : ${USE_STATIC_MKL}") @@ -177,7 +178,11 @@ function(caffe2_print_configuration_summary) if(${USE_NCCL}) message(STATUS " USE_SYSTEM_NCCL : ${USE_SYSTEM_NCCL}") endif() - message(STATUS " Found NVSHMEM : ${NVSHMEM_INCLUDE_DIR}") + if(${USE_ROCM}) + message(STATUS " Found ROCSHMEM : ${ROCSHMEM_INCLUDE_DIR}") + else() + message(STATUS " Found NVSHMEM : ${NVSHMEM_INCLUDE_DIR}") + endif() message(STATUS " USE_NNPACK : ${USE_NNPACK}") message(STATUS " USE_NUMPY : ${USE_NUMPY}") message(STATUS " USE_OBSERVERS : ${USE_OBSERVERS}") diff --git a/cmake/TorchConfig.cmake.in b/cmake/TorchConfig.cmake.in index abf5c8149116f..5749c67db112b 100644 --- a/cmake/TorchConfig.cmake.in +++ b/cmake/TorchConfig.cmake.in @@ -155,7 +155,7 @@ if(NOT @BUILD_SHARED_LIBS@) endif() set_target_properties(torch PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${TORCH_INCLUDE_DIRS}" - CXX_STANDARD 17 + CXX_STANDARD 20 ) if(TORCH_CXX_FLAGS) set_property(TARGET torch PROPERTY INTERFACE_COMPILE_OPTIONS "${TORCH_CXX_FLAGS}") diff --git a/cmake/public/LoadHIP.cmake b/cmake/public/LoadHIP.cmake index fe5621559b098..a87b16f5ba889 100644 --- a/cmake/public/LoadHIP.cmake +++ b/cmake/public/LoadHIP.cmake @@ -100,6 +100,9 @@ endmacro() # MODULE argument is added for clarity that CMake is searching # for FindHIP.cmake in Module mode find_package_and_print_version(HIP 1.0 MODULE) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + enable_language(HIP) +endif() if(HIP_FOUND) set(PYTORCH_FOUND_HIP TRUE) @@ -205,6 +208,7 @@ if(HIP_FOUND) find_package_and_print_version(rocthrust REQUIRED) find_package_and_print_version(hipsolver REQUIRED) find_package_and_print_version(rocsolver REQUIRED) + find_package_and_print_version(rocshmem) # workaround cmake 4 build issue if(CMAKE_VERSION VERSION_GREATER_EQUAL "4.0.0") message(WARNING "Work around hiprtc cmake failure for cmake >= 4") @@ -219,6 +223,7 @@ if(HIP_FOUND) if(UNIX) find_package_and_print_version(rccl) find_package_and_print_version(hsa-runtime64 REQUIRED) + find_package_and_print_version(rocm_smi REQUIRED) endif() # Optional components. diff --git a/cmake/public/cuda.cmake b/cmake/public/cuda.cmake index cb0e4e07c7f65..239641d352c41 100644 --- a/cmake/public/cuda.cmake +++ b/cmake/public/cuda.cmake @@ -78,8 +78,8 @@ endif() message(STATUS "PyTorch: CUDA detected: " ${CUDA_VERSION}) message(STATUS "PyTorch: CUDA nvcc is: " ${CUDA_NVCC_EXECUTABLE}) message(STATUS "PyTorch: CUDA toolkit directory: " ${CUDA_TOOLKIT_ROOT_DIR}) -if(CUDA_VERSION VERSION_LESS 12.0) - message(FATAL_ERROR "PyTorch requires CUDA 12.0 or above.") +if(CUDA_VERSION VERSION_LESS 12.1) + message(FATAL_ERROR "PyTorch requires CUDA 12.1 or above.") endif() if(CUDA_FOUND) @@ -307,6 +307,14 @@ else() endif() # nvrtc +# cuDNN frontend needs libnvrtc symbols, but linking through CUDA::nvrtc pulls +# CUDA::cuda_driver transitively. Keep a driver-free target for cuDNN users and +# reserve caffe2::nvrtc for the stub library that actually needs the driver API. +add_library(caffe2::nvrtc_runtime INTERFACE IMPORTED) +set_property( + TARGET caffe2::nvrtc_runtime PROPERTY INTERFACE_LINK_LIBRARIES + "${CUDA_NVRTC_LIB}") + add_library(caffe2::nvrtc INTERFACE IMPORTED) set_property( TARGET caffe2::nvrtc PROPERTY INTERFACE_LINK_LIBRARIES diff --git a/cmake/public/utils.cmake b/cmake/public/utils.cmake index 731bbe7d21862..a667ad55771b4 100644 --- a/cmake/public/utils.cmake +++ b/cmake/public/utils.cmake @@ -337,7 +337,7 @@ endmacro() # Usage: # torch_compile_options(lib_name) function(torch_compile_options libname) - set_property(TARGET ${libname} PROPERTY CXX_STANDARD 17) + set_property(TARGET ${libname} PROPERTY CXX_STANDARD 20) # until they can be unified, keep these lists synced with setup.py if(MSVC) @@ -391,12 +391,17 @@ function(torch_compile_options libname) endif() if(CMAKE_CXX_COMPILER_ID MATCHES "Clang") if(NOT USE_CUDA) - # NS: One can compile CUDA code with extra-semi flag as nvcc generates code like + # NS: One can not compile CUDA code with extra-semi flag as nvcc generates code like # namespace MemoryOps_cu_d8602b38_109889 __attribute__((visibility("hidden"))) { }; list(APPEND private_compile_options -Wextra-semi) else() # NVCC + clang15 reports deprecated copies from GPU lambda instantiations list(APPEND private_compile_options -Wno-deprecated-copy) + # NVCC + clang18 reports spurious deprecated deprecated literal operator declaration when there were none + # I.e. failures look like torch/headeronly/util/complex.h:334:40: error: identifier '_if' preceded by whitespace in a literal operator declaration is deprecated + # but if one to look at the source code, there are no space there + list(APPEND private_compile_options -Wno-deprecated-literal-operator) + endif() list(APPEND private_compile_options -Wmove) else() diff --git a/cmake/public/xpu.cmake b/cmake/public/xpu.cmake index b39e31d0ade8a..2731c2842c7f4 100644 --- a/cmake/public/xpu.cmake +++ b/cmake/public/xpu.cmake @@ -37,6 +37,11 @@ torch_xpu_get_arch_list(XPU_ARCH_FLAGS) # propagate to torch-xpu-ops set(TORCH_XPU_ARCH_LIST ${XPU_ARCH_FLAGS}) +# Ensure SYCL device code compiles with C++20 (matching CMAKE_CXX_STANDARD). +# SYCL_FLAGS flows into SYCL_COMPILE_FLAGS in torch-xpu-ops' BuildFlags.cmake +# and is passed directly to icpx on the device compilation command line. +list(APPEND SYCL_FLAGS -std=c++20) + # Ensure USE_XPU is enabled. string(APPEND XPU_HOST_CXX_FLAGS " -DUSE_XPU") string(APPEND XPU_HOST_CXX_FLAGS " -DSYCL_COMPILER_VERSION=${SYCL_COMPILER_VERSION}") diff --git a/docs/cpp/Makefile b/docs/cpp/Makefile index e244432b9fdb3..21379d1856690 100644 --- a/docs/cpp/Makefile +++ b/docs/cpp/Makefile @@ -2,24 +2,41 @@ # You can set these variables from the command line. SPHINXOPTS = -j auto -SPHINXBUILD = sphinx-build +SPHINXBUILD = python -m sphinx SPHINXPROJ = PyTorch SOURCEDIR = source BUILDDIR = build PYCMD = python +DOXYGEN = doxygen # Put it first so that "make" without argument is like "make help". help: @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) -.PHONY: help Makefile clean +.PHONY: help Makefile clean doxygen html + +# Generate Doxygen XML (required for Breathe directives) +doxygen: + @echo "Running Doxygen..." + @cd $(SOURCEDIR) && $(DOXYGEN) Doxyfile clean: - @# Clean up sphinx and doxygen build artifacts. + @# Clean up sphinx build artifacts and Doxygen output. + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + @rm -rf $(BUILDDIR)/xml + +# Build HTML with Doxygen XML generation first +html: doxygen @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) - @# Clean up exhale generated api. - @echo "Removing everything under generated 'source/api'..." - @rm -rf $(SOURCEDIR)/api + +# Check documentation coverage against curated allowlist and HTML formatting +coverage: html + @echo "Running C++ docs coverage check..." + @$(PYCMD) check_coverage.py + +coverage-only: + @echo "Running C++ docs coverage check (no rebuild)..." + @$(PYCMD) check_coverage.py # Catch-all target: route all unknown targets to Sphinx using the new # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). diff --git a/docs/cpp/README.md b/docs/cpp/README.md new file mode 100644 index 0000000000000..0808a5cb4af0b --- /dev/null +++ b/docs/cpp/README.md @@ -0,0 +1,315 @@ +# PyTorch C++ API Documentation + +This directory contains the source for PyTorch's C++ API documentation, built with +Sphinx, Breathe, and Doxygen. + +## How it works + +The documentation pipeline has three stages: + +``` +C++ headers ──→ Doxygen ──→ XML ──→ Breathe ──→ Sphinx ──→ HTML + ↑ ↑ + Doxyfile Markdown files with + (which headers) Breathe directives (MyST) +``` + +1. **Doxygen** reads C++ headers listed in `source/Doxyfile` `INPUT` and produces + XML in `build/xml/`. +2. **Breathe** is a Sphinx extension that reads Doxygen XML and makes it available + via directives like `` ```{doxygenclass} ClassName `` (MyST Markdown syntax). +3. **Sphinx** builds the final HTML from `.md` files in `source/`, using Breathe + directives to pull in C++ API documentation. The `myst_parser` extension + enables Markdown support. + +Only headers listed in the Doxyfile's `INPUT` are processed. Source files must +explicitly reference each symbol — nothing is auto-generated. + +## Building + +```bash +make html # Build Doxygen XML + Sphinx HTML +make doxygen # Build Doxygen XML only +make clean # Clean build artifacts +``` + +The output is in `build/html/`. + +## Contributing to the C++ docs + +### Adding a new API + +1. **Ensure the header is in the Doxyfile** — check that `source/Doxyfile` `INPUT` + includes the header file or its parent directory. If not, add it: + + ``` + INPUT = ... \ + ../../../path/to/your/header.h + ``` + +2. **Add a Breathe directive to the appropriate `.md` file** under `source/api/`. + See [Which directive to use](#which-directive-to-use) below. + +3. **Run `make html`** and check the output in `build/html/`. + +4. **Run `python check_coverage.py`** to verify your API shows as documented. + +### Which directive to use + +Use Breathe directives to pull documentation from Doxygen XML. These render the +full C++ signature, doc comments, parameters, and members automatically. + +All source files use MyST Markdown syntax (fenced directives with backticks). + +**Classes and structs** — use `doxygenclass` or `doxygenstruct`: + +````markdown +```{doxygenclass} torch::nn::Linear +:members: +:undoc-members: +``` +```` + +- `:members:` shows all public member functions and variables +- `:undoc-members:` includes members without doc comments +- Omit both flags to show only the class description (useful when `:members:` + causes rendering issues) + +**Free functions** — use `doxygenfunction`: + +````markdown +```{doxygenfunction} torch::autograd::grad +``` +```` + +For overloaded functions, Breathe will document all overloads. + +**Macros** — use `doxygendefine`: + +````markdown +```{doxygendefine} TORCH_LIBRARY +``` +```` + +**Typedefs** — use `doxygentypedef`: + +````markdown +```{doxygentypedef} torch::DeviceType +``` +```` + +### When Breathe directives don't work + +Some symbols can't be documented with Breathe directives: + +- **TORCH_MODULE holder classes** (e.g., `Conv2d`, `Linear`): The `TORCH_MODULE()` + macro generates these, but Doxygen can't index them. Document the `*Impl` class + instead (e.g., `Conv2dImpl`) — it contains all the actual methods. + +- **Functions with broken `\rst`/`\endrst` blocks**: Some doc comments use Doxygen's + `\rst` alias to embed RST. When a comment has multiple such blocks, Doxygen + generates malformed XML and Breathe renders raw text. In these cases, either: + - Fix the header to use native Doxygen (`@code{.cpp}`/`@endcode`, `@note`, + `@warning`) instead of `\rst`/`\endrst` + - Use a hand-written Sphinx C++ domain directive as a fallback + +- **Functions with mismatched `\param` names**: If a header's `\param` names don't + match the actual parameter names, Doxygen may fail to index the function. Use a + hand-written `cpp:function` directive instead. + +**Hand-written Sphinx C++ domain directives** (fallback): + +````markdown +```{cpp:function} void torch::autograd::backward(const variable_list& tensors, const variable_list& grad_tensors = {}, std::optional retain_graph = std::nullopt, bool create_graph = false, const variable_list& inputs = {}) + +Computes gradients of given tensors with respect to graph leaves. + +:param tensors: Tensors of which the derivative will be computed. +:param grad_tensors: The "vector" in the Jacobian-vector product. +``` +```` + +These don't pull from Doxygen — you write the signature and docs manually. + +**`{eval-rst}` escape hatch** — if a MyST directive doesn't render correctly, +you can embed raw RST: + +````markdown +```{eval-rst} +.. doxygenclass:: X::A + :members: + :protected-members: + :private-members: +``` +```` + +### Source file structure + +Each `.md` file under `source/api/` documents one topic area using MyST Markdown. +The typical pattern: + +````markdown +# Page Title + +Brief description of this API area. + +## Section Name + +Optional prose explaining usage, with a code example: + +```cpp +#include +auto x = torch::randn({2, 3}); +``` + +```{doxygenclass} torch::nn::SomeClass +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::SomeClassOptions +:members: +:undoc-members: +``` +```` + +**Nesting directives:** When a directive contains other directives or code blocks, +the outer fence must use more backticks than the inner ones: + +`````markdown +````{cpp:class} at::Tensor + +The primary tensor class. + +```{cpp:function} int64_t dim() const + +Returns the number of dimensions. +``` +```` +````` + +### Writing doc comments in C++ headers + +Doxygen extracts documentation from comments in headers. Use `///` style: + +```cpp +/// Applies a linear transformation to the incoming data: :math:`y = xA^T + b`. +/// +/// @code{.cpp} +/// auto linear = torch::nn::Linear(torch::nn::LinearOptions(10, 5)); +/// auto output = linear->forward(input); +/// @endcode +/// +/// @note The weight matrix is transposed compared to the Python API. +class LinearImpl : public Cloneable { +``` + +**Preferred Doxygen commands:** +- `@code{.cpp}` / `@endcode` for code examples +- `@note` for important notes +- `@warning` for warnings +- `@param name` for parameter descriptions +- `@return` for return value descriptions + +**Avoid** `\rst` / `\endrst` blocks — they cause rendering issues when a single +doc comment contains multiple blocks. Use native Doxygen commands instead. + +## Coverage checking + +```bash +make coverage # Build docs + run coverage check +make coverage-only # Run coverage check without rebuilding +``` + +`check_coverage.py` auto-discovers public APIs from Doxygen XML (`build/xml/index.xml`) +and checks which ones have Breathe or Sphinx directives in the source files. + +Reports are written to `cpp_coverage.txt` and `cpp_html_issues.txt`. + +### How coverage discovery works + +The script parses `build/xml/index.xml` to find all classes, structs, functions, +and macros that Doxygen indexed. It then: + +1. Filters out internal symbols using `EXCLUDED_PATTERNS` and `EXCLUDED_SYMBOLS` +2. Re-includes any symbols in `INCLUDED_SYMBOLS` (overrides exclusions) +3. Scans `.md` files for Breathe (`doxygenclass`, `doxygenfunction`, etc.) and + Sphinx C++ domain (`cpp:class`, `cpp:function`, etc.) directives +4. Reports the gap + +### Excluding internal APIs from coverage + +Not everything in Doxygen XML is public API. The script has three exclusion mechanisms: + +**`EXCLUDED_PATTERNS`** — regex patterns for broad categories: + +```python +EXCLUDED_PATTERNS = [ + r".*::detail::.*", # Internal namespaces + r"torch::jit::.*", # Deprecated + r".*::_\w+", # Underscore-prefixed internals +] +``` + +**`EXCLUDED_SYMBOLS`** — exact match for specific symbols: + +```python +EXCLUDED_SYMBOLS = { + "torch::autograd::deleteNode", + "at::native::dataSize", +} +``` + +**Doxyfile `EXCLUDE`** — prevents Doxygen from indexing files entirely: + +``` +EXCLUDE = ../../../torch/csrc/api/include/torch/detail +``` + +### Including "internal" APIs that are actually public + +Some APIs live in internal-looking namespaces but are widely used by external +developers. To track these for documentation coverage, add them to +`INCLUDED_SYMBOLS` in `check_coverage.py`: + +```python +INCLUDED_SYMBOLS: set[str] = { + "c10::IValue", # Used for custom op registration +} +``` + +`INCLUDED_SYMBOLS` takes priority over all exclusions. + +### HTML formatting checks + +The script also checks built HTML for: +- Unresolved Breathe directives ("Cannot find class/struct/function") +- Raw directive text in output (build failures) +- Sphinx "problematic" nodes (broken references) +- Near-empty API pages + +### coverxygen (optional) + +```bash +python check_coverage.py --coverxygen +``` + +This runs [coverxygen](https://github.com/psycofdj/coverxygen) on the Doxygen XML +to measure what percentage of C++ symbols have doc comments in the source code. +This is complementary to file coverage — it tells you which headers need more +doc comments, not which symbols need Sphinx directives. + +## Known issues + +- **`\rst`/`\endrst` rendering** — Some C++ headers use Doxygen's `\rst` alias to + embed RST in doc comments. When a single doc comment has multiple `\rst`/`\endrst` + blocks, Doxygen generates malformed XML and Breathe renders them as raw text. + Avoid using `:members:` on classes with this issue. Use hand-written examples + in the source files instead, or convert the headers to native Doxygen commands. + +- **TORCH_MODULE holder classes** — Doxygen can't index classes generated by the + `TORCH_MODULE()` macro. Use the `*Impl` class name in directives instead. + +- **"Subclassed by" plain text** — Without per-class pages (exhale), Breathe + renders "Subclassed by" lists as plain text instead of links. A `conf.py` hook + (`remove_subclassed_by`) strips these automatically. diff --git a/docs/cpp/check_coverage.py b/docs/cpp/check_coverage.py new file mode 100644 index 0000000000000..bd958671cacbf --- /dev/null +++ b/docs/cpp/check_coverage.py @@ -0,0 +1,782 @@ +#!/usr/bin/env python3 +"""C++ API documentation coverage checker. + +Auto-discovers public C++ APIs from Doxygen XML output and checks which ones +are documented in the RST source files via Breathe or Sphinx C++ domain +directives. + +Uses an exclusion list (EXCLUDED_APIS) to skip internal/detail symbols that +don't need public documentation, rather than maintaining a hardcoded allowlist. + +Additionally checks built HTML for broken formatting (empty pages, +unresolved directives, rendering errors). + +Usage: + python check_coverage.py # RST coverage + HTML checks + python check_coverage.py --coverxygen # also run coverxygen on Doxygen XML +""" + +import argparse +import re +import subprocess +import sys +import xml.etree.ElementTree as ET +from pathlib import Path + + +# ─── Paths ─────────────────────────────────────────────────────────────────── + +SCRIPT_DIR = Path(__file__).resolve().parent +SOURCE_DIR = SCRIPT_DIR / "source" +BUILD_HTML = SCRIPT_DIR / "build" / "html" +BUILD_XML = SCRIPT_DIR / "build" / "xml" +COVERAGE_OUTPUT = SCRIPT_DIR / "cpp_coverage.txt" +HTML_REPORT = SCRIPT_DIR / "cpp_html_issues.txt" + +# ─── Inclusion override ────────────────────────────────────────────────────── +# Symbols that match an exclusion pattern but should still be tracked. +# Use this for "internal" APIs that are widely used as public API. +# Add the fully-qualified symbol name here and it will bypass all exclusions. +INCLUDED_SYMBOLS: set[str] = { + # Example: "c10::IValue" would track it even though c10::IValue is excluded +} + +# ─── Exclusion list ────────────────────────────────────────────────────────── +# Symbols that should NOT be flagged as missing documentation. +# Add internal, detail, or otherwise non-public symbols here. +# Note: INCLUDED_SYMBOLS takes priority over these exclusions. + +EXCLUDED_PATTERNS = [ + # Internal/detail namespaces + r".*::detail::.*", + r".*::detail_::.*", + r"torch::python::.*", + # Underscore-prefixed internal classes + r".*::_\w+", + # Enum helper structs + r"torch::enumtype::.*", + # OptimizerCloneableOptions SFINAE helpers + r"torch::optim::OptimizerCloneableOptions::.*", + # Internal optimizer state/options cloneable helpers + r"torch::optim::OptimizerCloneable.*", + # Error classes (c10 exceptions) + r"c10::.*Error$", + r"c10::ErrorAlwaysShowCppStacktrace", + # Warning internals + r"c10::Warning.*", + r"c10::WarningHandler", + r"c10::WarningUtils::.*", + # c10 IValue internals + r"c10::IValue::.*", + r"c10::IValue", + r"c10::WeakIValue", + r"c10::ivalue::.*", + r"c10::StrongTypePtr", + r"c10::WeakTypePtr", + r"c10::WeakOrStrongTypePtr", + r"c10::WeakOrStrongCompilationUnit", + r"c10::Capsule", + r"c10::OptionalArray", + r"c10::StreamData3", + # OrderedDict::Item (internal helper) + r"torch::OrderedDict::Item", + # ExpandingArray (internal template utility) + r"torch::ExpandingArray.*", + # IMethod (internal) + r"torch::IMethod", + # CustomClassHolder (internal base) + r"torch::CustomClassHolder", + # NodeGuard (internal autograd) + r"torch::autograd::NodeGuard", + # Autograd internals + r"torch::autograd::CppNode", + r"torch::autograd::ExtractVariables", + r"torch::autograd::Node", + r"torch::autograd::Node::.*", + r"torch::autograd::TraceableFunction", + r"torch::autograd::TypeAndSize", + # Sequencer internals + r"torch::data::.*::detail::.*", + # cuDNN descriptor internals + r"at::native::ActivationDescriptor", + r"at::native::ConvolutionDescriptor", + r"at::native::SpatialTransformerDescriptor", + r"at::native::DropoutDescriptor", + r"at::native::RNNDataDescriptor", + r"at::native::DftiDescriptor", + r"at::native::DescriptorDeleter", + r"at::native::DftiDescriptorDeleter", + r"at::native::RNNDescriptor", + # ATen internals + r"at::OptionalTensorRef", + r"at::TensorRef", + # at::cuda internals (allocator, workspace, cublas) + r"at::cuda::WorkspaceMapWithMutex", + r"at::cuda::clearCublasWorkspaces.*", + r"at::cuda::cublas_handle_stream_to_workspace", + r"at::cuda::cublaslt_handle_stream_to_workspace", + r"at::cuda::getCUDABlasLt.*", + r"at::cuda::getCUDADeviceAllocator", + r"at::cuda::getChosenWorkspaceSize", + r"at::cuda::getNumGPUs", + r"at::cuda::is_available", + r"at::cuda::warp_size", + # jit namespace (deprecated) + r"torch::jit::.*", + # Operators that are just operator<< or operator>> + r".*::operator<<", + r".*::operator>>", + r".*::operator==", + r".*::operator!=", + # Internal serialize helpers + r"torch::optim::serialize", + r"torch::optim::detail::.*", + # Reduction enum helpers + r"torch::nn::reduction", + r"torch::nn::log_target", + # Internal module utils + r"torch::nn::modules::utils::.*", + # Internal c10 helpers + r"c10::detail::.*", + r"c10::detail_::.*", + r"c10::makeArrayRef", + r"c10::checkObjectSortSchema", + r"c10::getGreaterThanComparator", + r"c10::getLessThanComparator", + r"c10::value_or_else", + r"c10::warn", + r"c10::GetExceptionString", + # torch::detail + r"torch::detail::.*", + # Internal data shuttle/queue + r"torch::data::detail::.*", + # DataLoaderBase internal types + r"torch::data::DataLoaderBase::.*", + r"torch::data::WorkerException", + r"torch::data::FullDataLoaderOptions", + # Template specializations of Stack + r"torch::data::transforms::Stack< .*>", + # Example partial specialization + r"torch::data::Example< .*>", + # Doxygen internal macros + r"DEFINE_CASE", + r"DEFINE_TAG", + r"COUNT_TAG", + r"TRUTH_TABLE_ENTRY", + r"C10_EXPAND_MSVC_WORKAROUND", + r"TORCH_FORALL_TAGS", + # Non-public torch::nn functions (module stream operators, etc.) + r"torch::nn::operator.*", + # AnyModule/AnyValue internal holders + r"torch::nn::AnyModuleHolder.*", + r"torch::nn::AnyModulePlaceholder", + r"torch::nn::AnyValue.*", + r"torch::nn::NamedAnyModule", + # Internal base classes (users use the derived classes) + r"torch::nn::ConvNdImpl", + r"torch::nn::ConvTransposeNdImpl", + r"torch::nn::BatchNormImplBase", + r"torch::nn::NormImplBase", + r"torch::nn::InstanceNormImpl", + r"torch::nn::MaxPoolImpl", + r"torch::nn::AvgPoolImpl", + r"torch::nn::AdaptiveAvgPoolImpl", + r"torch::nn::AdaptiveMaxPoolImpl", + r"torch::nn::MaxUnpoolImpl", + r"torch::nn::LPPoolImpl", + r"torch::nn::ConstantPadImpl", + r"torch::nn::ReflectionPadImpl", + r"torch::nn::ReplicationPadImpl", + r"torch::nn::ZeroPadImpl", + r"torch::nn::FractionalMaxPoolImpl", + # nn::functions internal namespace + r"torch::nn::functions::.*", + # AdaptiveLogSoftmaxWithLoss (niche, rarely used in C++) + r"torch::nn::AdaptiveLogSoftmaxWithLoss.*", + r"torch::nn::ASMoutput", + # CrossMapLRN2d (niche) + r"torch::nn::CrossMapLRN2d.*", + # _out function variants (documented alongside the main function) + r"torch::special::.*_out", + r"torch::fft::.*_out", + # torch internal helpers + r"torch::InitLambda", + r"torch::dispatch", + r"torch::equal_if_defined", + r"torch::getAllCustomClassesNames", + r"torch::init", + r"torch::make_custom_class", + r"torch::selective_class_", + r"torch::pickle_load", + r"torch::pickle_save", + r"torch::schema", + r"torch::nativert::.*", + # RNNCellOptionsBase (internal base) + r".*::RNNCellOptionsBase", + # Unnamespaced Options structs (indexed without namespace by Doxygen) + r"^[A-Z]\w+Options$", + # Unnamespaced classes without namespace (Doxygen quirk) + r"^TransformerDecoderLayer$", + r"^TransformerDecoderLayerOptions$", + # functional namespace internal options structs + r"functional::.*FuncOptions", +] + +# Specific symbols to exclude (exact match) +EXCLUDED_SYMBOLS = { + # Internal / not useful to document individually + "torch::data::datasets::map", + "torch::data::datasets::make_shared_dataset", + "torch::data::datasets::operator<<", + "torch::data::datasets::operator>>", + "torch::enumtype::get_enum_name", + "torch::enumtype::reduction_get_enum", + "torch::autograd::_wrap_outputs", + "torch::autograd::check_variable_result", + "torch::autograd::CppNode_apply_functional", + "torch::autograd::CppNode_apply_functional_ivalue", + "torch::autograd::forward_ad::enter_dual_level", + "torch::autograd::forward_ad::exit_dual_level", + "torch::autograd::any_variable_requires_grad", + "torch::autograd::collect_next_edges", + "torch::autograd::create_gradient_edge", + "torch::autograd::deleteNode", + "torch::autograd::extract_vars", + "torch::autograd::get_current_node", + "torch::autograd::to_optional", + "torch::autograd::to_output_type", + "torch::nn::parallel::replicate", + "torch::nn::parallel::parallel_apply", + "torch::nn::parallel::data_parallel", + "torch::python::add_module_bindings", + "torch::python::bind_module", + "torch::python::init_bindings", + # at::native cuDNN internals + "at::native::dataSize", + "at::native::fixSizeOneDimStride", + "at::native::operator<<", + "at::native::getCudnnDataTypeFromScalarType", + # c10 cuda pool functions (internal) + "c10::cuda::getStreamFromPool", + "c10::cuda::getStreamFromExternal", + "c10::xpu::getStreamFromPool", + "c10::xpu::getStreamFromExternal", + # c10 private use backend registration (internal) + "c10::get_privateuse1_backend", + "c10::is_privateuse1_backend_registered", + "c10::register_privateuse1_backend", + "c10::isValidDeviceType", + "c10::DeviceTypeName", + # torch::stable::detail internals + "torch::stable::detail::unbox_to_tuple_impl", + "torch::stable::detail::unbox_to_tuple", + "torch::stable::detail::box_from_tuple_impl", + "torch::stable::detail::box_from_tuple", + # torch::stable::accelerator (documented in stable API page) + "torch::stable::accelerator::getCurrentStream", +} + +# Namespaces whose free functions should be checked for documentation +PUBLIC_FUNCTION_NAMESPACES = { + "torch", + "torch::autograd", + "torch::cuda", + "torch::mps", + "torch::xpu", + "torch::fft", + "torch::special", + "torch::nn::functional", + "torch::nn::init", + "torch::nn::utils", + "torch::nn::utils::rnn", + "torch::data", + "torch::stable", + "torch::stable::accelerator", + "c10", + "c10::cuda", + "c10::xpu", + "at::cuda", +} + + +# ─── XML parsing ───────────────────────────────────────────────────────────── + + +def _is_excluded(symbol: str) -> bool: + """Check if a symbol should be excluded from coverage tracking.""" + if symbol in INCLUDED_SYMBOLS: + return False + if symbol in EXCLUDED_SYMBOLS: + return True + for pattern in EXCLUDED_PATTERNS: + if re.fullmatch(pattern, symbol): + return True + return False + + +def _categorize(name: str) -> str: + """Assign a category based on the symbol's namespace.""" + if name.startswith("torch::nn::functional::"): + return "torch::nn::functional" + if name.startswith("torch::nn::init::"): + return "torch::nn::init" + if name.startswith("torch::nn::utils::"): + return "torch::nn::utils" + if name.startswith("torch::nn::"): + # Distinguish modules from other nn symbols + short = name.split("::")[-1] + if short[0].isupper(): + return "torch::nn (modules)" + return "torch::nn" + if name.startswith("torch::optim::"): + return "torch::optim" + if name.startswith("torch::data::"): + return "torch::data" + if name.startswith("torch::autograd::"): + return "torch::autograd" + if name.startswith("torch::serialize::") or name in ("torch::save", "torch::load"): + return "torch::serialize" + if name.startswith("torch::stable::"): + return "torch::stable" + if name.startswith("torch::fft::"): + return "torch::fft" + if name.startswith("torch::special::"): + return "torch::special" + if name.startswith(("torch::cuda::", "torch::mps::", "torch::xpu::")): + return "torch (device)" + if name.startswith("torch::"): + return "torch (core)" + if name.startswith("c10::cuda::"): + return "c10::cuda" + if name.startswith("c10::xpu::"): + return "c10::xpu" + if name.startswith("c10::"): + return "c10" + if name.startswith("at::cuda::"): + return "at::cuda" + if name.startswith("at::"): + return "at" + return "other" + + +def discover_apis_from_xml(xml_dir: Path) -> dict[str, list[tuple[str, str]]]: + """Parse Doxygen index.xml to discover all public APIs. + + Returns dict of category -> list of (symbol, kind). + """ + index_path = xml_dir / "index.xml" + if not index_path.exists(): + print( + f"ERROR: {index_path} not found. Run 'make doxygen' first.", + file=sys.stderr, + ) + sys.exit(1) + + tree = ET.parse(index_path) + root = tree.getroot() + + apis: dict[str, list[tuple[str, str]]] = {} + + # Collect classes and structs + for compound in root.findall("compound"): + kind = compound.get("kind") + if kind not in ("class", "struct"): + continue + name = compound.find("name").text + if _is_excluded(name): + continue + category = _categorize(name) + apis.setdefault(category, []).append((name, kind)) + + # Collect free functions from public namespaces + for compound in root.findall("compound"): + if compound.get("kind") != "namespace": + continue + ns_name = compound.find("name").text + if ns_name not in PUBLIC_FUNCTION_NAMESPACES: + continue + seen_funcs = set() + for member in compound.findall("member"): + if member.get("kind") != "function": + continue + func_name = member.find("name").text + qualified = f"{ns_name}::{func_name}" + if qualified in seen_funcs: + continue # skip overloads + seen_funcs.add(qualified) + if _is_excluded(qualified): + continue + category = _categorize(qualified) + apis.setdefault(category, []).append((qualified, "function")) + + # Collect macros (defines) from file compounds + for compound in root.findall("compound"): + if compound.get("kind") != "file": + continue + for member in compound.findall("member"): + if member.get("kind") != "define": + continue + macro_name = member.find("name").text + # Only track well-known public macros + if macro_name.startswith(("TORCH_LIBRARY", "TORCH_MODULE")): + if _is_excluded(macro_name): + continue + apis.setdefault("torch (macros)", []).append((macro_name, "define")) + + # Sort each category and deduplicate + for category in apis: + apis[category] = sorted(set(apis[category])) + + return apis + + +# ─── Source scanning ───────────────────────────────────────────────────────── + +# RST directives: .. doxygenclass:: torch::nn::ReLU +RST_DIRECTIVE_RE = re.compile( + r"^\.\.\s+doxygen(class|struct|function|typedef|define|enum|namespace)" + r"::\s*(.+?)\s*$", + re.MULTILINE, +) + +RST_CPP_DIRECTIVE_RE = re.compile( + r"^\.\.\s+cpp:(class|struct|function|enum|type)" r"::\s*(.+?)\s*$", + re.MULTILINE, +) + +# MyST directives: ```{doxygenclass} torch::nn::ReLU +MYST_DIRECTIVE_RE = re.compile( + r"^`{3,}\{doxygen(class|struct|function|typedef|define|enum|namespace)\}\s*(.+?)\s*$", + re.MULTILINE, +) + +MYST_CPP_DIRECTIVE_RE = re.compile( + r"^`{3,}\{cpp:(class|struct|function|enum|type)\}\s*(.+?)\s*$", + re.MULTILINE, +) + + +def scan_sources(source_dir: Path) -> set[str]: + """Extract all documented symbols from RST/MyST breathe and cpp domain directives.""" + documented = set() + for src_file in list(source_dir.rglob("*.rst")) + list(source_dir.rglob("*.md")): + content = src_file.read_text(errors="replace") + patterns = ( + RST_DIRECTIVE_RE, + RST_CPP_DIRECTIVE_RE, + MYST_DIRECTIVE_RE, + MYST_CPP_DIRECTIVE_RE, + ) + for pattern in patterns: + for match in pattern.finditer(content): + symbol = match.group(2) + # Strip template prefix + if symbol.startswith("template"): + gt = symbol.find(">") + if gt != -1: + symbol = symbol[gt + 1 :].lstrip() + # Strip function signature + paren = symbol.find("(") + if paren != -1: + symbol = symbol[:paren].rstrip() + documented.add(symbol) + return documented + + +# ─── Coverage report ───────────────────────────────────────────────────────── + + +def generate_coverage_report( + apis: dict[str, list[tuple[str, str]]], documented: set[str] +) -> str: + """Generate a coverage report comparing discovered APIs against RST docs.""" + lines = [] + lines.append("Undocumented C++ objects") + lines.append("=" * 50) + lines.append("") + + total = 0 + total_missing = 0 + section_stats = [] + + for category in sorted(apis.keys()): + symbols = apis[category] + section_missing = [] + for symbol, kind in symbols: + total += 1 + unqualified = symbol.rsplit("::", 1)[-1] + if symbol not in documented and unqualified not in documented: + section_missing.append((symbol, kind)) + total_missing += 1 + + covered = len(symbols) - len(section_missing) + section_stats.append((category, covered, len(symbols))) + + if section_missing: + lines.append(category) + lines.append("-" * len(category)) + for symbol, kind in section_missing: + lines.append(f" * {symbol} ({kind})") + lines.append("") + + # Summary + total_covered = total - total_missing + pct = (total_covered / total * 100) if total else 0 + + lines.append("") + lines.append("=" * 50) + lines.append("Summary") + lines.append("=" * 50) + lines.append("") + lines.append(f"Total APIs discovered: {total}") + lines.append(f"Documented: {total_covered}") + lines.append(f"Missing: {total_missing}") + lines.append(f"Coverage: {pct:.1f}%") + lines.append("") + + # Per-section table + lines.append(f"{'Category':<45} {'Covered':>8} {'Total':>6} {'%':>7}") + lines.append("-" * 70) + for category, covered, section_total in section_stats: + spct = (covered / section_total * 100) if section_total else 0 + lines.append(f"{category:<45} {covered:>8} {section_total:>6} {spct:>6.1f}%") + lines.append("") + + return "\n".join(lines) + + +# ─── HTML checks ───────────────────────────────────────────────────────────── + +BROKEN_PATTERNS = [ + ( + re.compile(r"Cannot find (?:class|struct|function|file)", re.IGNORECASE), + "unresolved breathe directive", + ), + ( + re.compile(r"Unable to resolve (?:class|struct|function)", re.IGNORECASE), + "unresolved breathe directive (ambiguous overload)", + ), + ( + re.compile(r"doxygenclass:|doxygenfunction:|doxygenstruct:", re.IGNORECASE), + "raw directive text in output", + ), + ( + re.compile(r"", re.IGNORECASE), + "Sphinx problematic node (broken reference)", + ), + ( + re.compile(r"System Message:", re.IGNORECASE), + "Sphinx system message (build error)", + ), +] + +MIN_CONTENT_LENGTH = 500 + + +def check_html_output(build_dir: Path) -> str: + """Check built HTML for broken formatting and empty pages.""" + issues = [] + + if not build_dir.exists(): + return "ERROR: build/html directory not found. Run 'make html' first.\n" + + for html_file in sorted(build_dir.rglob("*.html")): + rel = html_file.relative_to(build_dir) + if rel.name in ("search.html", "genindex.html", "objects.inv"): + continue + + try: + content = html_file.read_text(errors="replace") + except Exception as e: + issues.append((str(rel), f"cannot read: {e}")) + continue + + for pattern, description in BROKEN_PATTERNS: + matches = pattern.findall(content) + if matches: + issues.append((str(rel), f"{description} ({len(matches)}x)")) + + if str(rel).startswith("api/"): + text = re.sub(r"<[^>]+>", "", content) + text = re.sub(r"\s+", " ", text).strip() + if len(text) < MIN_CONTENT_LENGTH: + issues.append((str(rel), f"possibly empty page ({len(text)} chars)")) + + lines = [] + lines.append("HTML Formatting Check") + lines.append("=" * 50) + lines.append("") + + if not issues: + lines.append("No issues found.") + else: + lines.append(f"Found {len(issues)} issue(s):") + lines.append("") + lines.append(f"{'File':<55} Issue") + lines.append("-" * 90) + for filepath, issue in issues: + lines.append(f"{filepath:<55} {issue}") + + lines.append("") + return "\n".join(lines) + + +# ─── coverxygen integration ───────────────────────────────────────────────── + + +def run_coverxygen(xml_dir: Path) -> str: + """Run coverxygen on Doxygen XML output for doc-comment coverage.""" + lines = [] + lines.append("Coverxygen Report (Doxygen doc-comment coverage)") + lines.append("=" * 50) + lines.append("") + + if not xml_dir.exists(): + lines.append("ERROR: build/xml directory not found. Run 'make doxygen' first.") + return "\n".join(lines) + + coverxygen_cmd = None + for cmd in [ + ["coverxygen", "--version"], + [sys.executable, "-m", "coverxygen", "--version"], + ]: + try: + subprocess.run(cmd, capture_output=True, check=True) + coverxygen_cmd = cmd[:-1] + break + except (FileNotFoundError, subprocess.CalledProcessError): + continue + if coverxygen_cmd is None: + lines.append("coverxygen not installed. Install with: pip install coverxygen") + lines.append("") + lines.append("Once installed, coverxygen analyzes Doxygen XML to report what") + lines.append("percentage of C++ symbols have doc comments in the source code.") + lines.append("This is complementary to the RST coverage check above.") + lines.append("") + lines.append("Usage:") + lines.append( + f" coverxygen --xml-dir {xml_dir} --src-dir ../../ --output coverxygen.info" + ) + lines.append(" # Then use lcov/genhtml to visualize:") + lines.append( + " genhtml --no-function-coverage coverxygen.info -o coverxygen_html" + ) + return "\n".join(lines) + + try: + result = subprocess.run( + coverxygen_cmd + + [ + "--xml-dir", + str(xml_dir), + "--src-dir", + str(SCRIPT_DIR / ".." / ".."), + "--output", + "-", + "--kind", + "class,struct,function", + "--scope", + "public", + "--exclude", + ".*/build/.*", + "--exclude", + ".*/detail/.*", + "--exclude", + ".*/nativert/.*", + "--exclude", + ".*/stable/library\\.h", + ], + capture_output=True, + text=True, + timeout=120, + ) + if result.returncode == 0: + total = 0 + documented_count = 0 + for line in result.stdout.splitlines(): + if line.startswith("DA:"): + total += 1 + parts = line.split(",") + if len(parts) >= 2 and parts[1].strip() != "0": + documented_count += 1 + pct = (documented_count / total * 100) if total else 0 + lines.append(f"Symbols scanned: {total}") + lines.append(f"With doc comments: {documented_count}") + lines.append(f"Coverage: {pct:.1f}%") + lines.append("") + lines.append("Full lcov output saved to: coverxygen.info") + (SCRIPT_DIR / "coverxygen.info").write_text(result.stdout) + else: + lines.append(f"coverxygen failed (exit {result.returncode}):") + lines.append(result.stderr[:500]) + except subprocess.TimeoutExpired: + lines.append("coverxygen timed out after 120s") + except Exception as e: + lines.append(f"coverxygen error: {e}") + + lines.append("") + return "\n".join(lines) + + +# ─── Main ──────────────────────────────────────────────────────────────────── + + +def main(): + parser = argparse.ArgumentParser(description="C++ docs coverage checker") + parser.add_argument( + "--coverxygen", + action="store_true", + help="Also run coverxygen on Doxygen XML for doc-comment coverage", + ) + parser.add_argument( + "--html-only", + action="store_true", + help="Only run HTML formatting checks", + ) + args = parser.parse_args() + + reports = [] + + if not args.html_only: + # Phase 1: Discover APIs from Doxygen XML + print("Discovering APIs from Doxygen XML...") + apis = discover_apis_from_xml(BUILD_XML) + total_apis = sum(len(v) for v in apis.values()) + print(f" Found {total_apis} public APIs across {len(apis)} categories") + + # Phase 2: Scan RST for documented symbols + print("Scanning sources for breathe directives...") + documented = scan_sources(SOURCE_DIR) + print(f" Found {len(documented)} documented symbols") + + coverage_report = generate_coverage_report(apis, documented) + reports.append(coverage_report) + + COVERAGE_OUTPUT.write_text(coverage_report) + print(f" Coverage report written to: {COVERAGE_OUTPUT}") + + # Phase 3: HTML checks + print("Checking HTML output for formatting issues...") + html_report = check_html_output(BUILD_HTML) + reports.append(html_report) + HTML_REPORT.write_text(html_report) + print(f" HTML report written to: {HTML_REPORT}") + + # Phase 4: coverxygen (optional) + if args.coverxygen: + print("Running coverxygen...") + cov_report = run_coverxygen(BUILD_XML) + reports.append(cov_report) + + # Print everything + print() + print("=" * 60) + for report in reports: + print(report) + print() + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/docs/cpp/cpp_coverage.txt b/docs/cpp/cpp_coverage.txt new file mode 100644 index 0000000000000..4fce09f90a646 --- /dev/null +++ b/docs/cpp/cpp_coverage.txt @@ -0,0 +1,292 @@ +Undocumented C++ objects +================================================== + +torch (device) +-------------- + * torch::cuda::device_count (function) + * torch::cuda::is_available (function) + * torch::cuda::manual_seed (function) + * torch::cuda::manual_seed_all (function) + * torch::cuda::synchronize (function) + * torch::mps::commit (function) + * torch::mps::get_command_buffer (function) + * torch::mps::get_dispatch_queue (function) + * torch::mps::is_available (function) + * torch::mps::manual_seed (function) + * torch::mps::synchronize (function) + +torch (macros) +-------------- + * TORCH_MODULE (define) + * TORCH_MODULE_IMPL (define) + +torch::data +----------- + * torch::data::datasets::ChunkDatasetOptions (struct) + * torch::data::datasets::TensorDataset (struct) + * torch::data::make_data_loader (function) + * torch::data::samplers::BatchSize (struct) + * torch::data::samplers::CustomBatchRequest (struct) + +torch::fft +---------- + * torch::fft::fft (function) + * torch::fft::fft2 (function) + * torch::fft::fftfreq (function) + * torch::fft::fftn (function) + * torch::fft::fftshift (function) + * torch::fft::hfft (function) + * torch::fft::hfft2 (function) + * torch::fft::hfftn (function) + * torch::fft::ifft (function) + * torch::fft::ifft2 (function) + * torch::fft::ifftn (function) + * torch::fft::ifftshift (function) + * torch::fft::ihfft (function) + * torch::fft::ihfft2 (function) + * torch::fft::ihfftn (function) + * torch::fft::irfft (function) + * torch::fft::irfft2 (function) + * torch::fft::irfftn (function) + * torch::fft::rfft (function) + * torch::fft::rfft2 (function) + * torch::fft::rfftfreq (function) + * torch::fft::rfftn (function) + +torch::nn (modules) +------------------- + * torch::nn::AdaptiveAvgPoolOptions (struct) + * torch::nn::AdaptiveMaxPoolOptions (struct) + * torch::nn::AvgPoolOptions (struct) + * torch::nn::BatchNormOptions (struct) + * torch::nn::BilinearOptions (struct) + * torch::nn::CELUOptions (struct) + * torch::nn::ConstantPad1dImpl (class) + * torch::nn::ConstantPad2dImpl (class) + * torch::nn::ConstantPad3dImpl (class) + * torch::nn::ConstantPadOptions (struct) + * torch::nn::ConvOptions (struct) + * torch::nn::ConvTransposeOptions (struct) + * torch::nn::CosineSimilarityImpl (class) + * torch::nn::CosineSimilarityOptions (struct) + * torch::nn::DropoutOptions (struct) + * torch::nn::ELUOptions (struct) + * torch::nn::EmbeddingBagFromPretrainedOptions (struct) + * torch::nn::EmbeddingBagOptions (struct) + * torch::nn::EmbeddingFromPretrainedOptions (struct) + * torch::nn::EmbeddingOptions (struct) + * torch::nn::FlattenOptions (struct) + * torch::nn::FoldImpl (class) + * torch::nn::FractionalMaxPoolOptions (struct) + * torch::nn::Functional (class) + * torch::nn::GELUOptions (struct) + * torch::nn::GLUOptions (struct) + * torch::nn::HardshrinkOptions (struct) + * torch::nn::HardtanhOptions (struct) + * torch::nn::InstanceNormOptions (struct) + * torch::nn::KLDivLossOptions (struct) + * torch::nn::L1LossOptions (struct) + * torch::nn::LPPoolOptions (struct) + * torch::nn::LayerNormOptions (struct) + * torch::nn::LeakyReLUOptions (struct) + * torch::nn::LinearOptions (struct) + * torch::nn::LogSoftmaxOptions (struct) + * torch::nn::MaxPoolOptions (struct) + * torch::nn::MaxUnpoolOptions (struct) + * torch::nn::ModuleDictImpl (class) + * torch::nn::ModuleListImpl (class) + * torch::nn::MultiheadAttentionOptions (struct) + * torch::nn::PReLUOptions (struct) + * torch::nn::PairwiseDistanceImpl (class) + * torch::nn::PairwiseDistanceOptions (struct) + * torch::nn::PixelShuffleImpl (struct) + * torch::nn::PixelUnshuffleImpl (struct) + * torch::nn::RNNOptions (struct) + * torch::nn::RReLUOptions (struct) + * torch::nn::ReLU6Options (struct) + * torch::nn::ReLUOptions (struct) + * torch::nn::ReflectionPad1dImpl (class) + * torch::nn::ReflectionPad2dImpl (class) + * torch::nn::ReflectionPad3dImpl (class) + * torch::nn::ReflectionPadOptions (struct) + * torch::nn::ReplicationPad1dImpl (class) + * torch::nn::ReplicationPad2dImpl (class) + * torch::nn::ReplicationPad3dImpl (class) + * torch::nn::ReplicationPadOptions (struct) + * torch::nn::SELUOptions (struct) + * torch::nn::SequentialImpl (class) + * torch::nn::SoftmaxOptions (struct) + * torch::nn::SoftminOptions (struct) + * torch::nn::SoftplusOptions (struct) + * torch::nn::SoftshrinkOptions (struct) + * torch::nn::ThresholdOptions (struct) + * torch::nn::TransformerDecoderOptions (struct) + * torch::nn::TransformerEncoderLayerOptions (struct) + * torch::nn::TransformerEncoderOptions (struct) + * torch::nn::TransformerOptions (struct) + * torch::nn::UnflattenOptions (struct) + * torch::nn::UnfoldImpl (class) + * torch::nn::UpsampleImpl (class) + * torch::nn::ZeroPad1dImpl (class) + * torch::nn::ZeroPad2dImpl (class) + * torch::nn::ZeroPad3dImpl (class) + * torch::nn::ZeroPadOptions (struct) + +torch::nn::functional +--------------------- + * torch::nn::functional::AlphaDropoutFuncOptions (struct) + * torch::nn::functional::BatchNormFuncOptions (struct) + * torch::nn::functional::ConvFuncOptions (struct) + * torch::nn::functional::ConvTransposeFuncOptions (struct) + * torch::nn::functional::DropoutFuncOptions (struct) + * torch::nn::functional::EmbeddingBagFuncOptions (struct) + * torch::nn::functional::EmbeddingFuncOptions (struct) + * torch::nn::functional::FeatureAlphaDropoutFuncOptions (struct) + * torch::nn::functional::GridSampleFuncOptions (struct) + * torch::nn::functional::GumbelSoftmaxFuncOptions (struct) + * torch::nn::functional::InstanceNormFuncOptions (struct) + * torch::nn::functional::InterpolateFuncOptions (struct) + * torch::nn::functional::LogSoftmaxFuncOptions (struct) + * torch::nn::functional::MaxUnpoolFuncOptions (struct) + * torch::nn::functional::MultiheadAttentionForwardFuncOptions (struct) + * torch::nn::functional::PadFuncOptions (struct) + * torch::nn::functional::RReLUFuncOptions (struct) + * torch::nn::functional::SoftmaxFuncOptions (struct) + * torch::nn::functional::SoftminFuncOptions (struct) + * torch::nn::functional::adaptive_max_pool2d_with_indices (function) + * torch::nn::functional::adaptive_max_pool3d_with_indices (function) + * torch::nn::functional::fractional_max_pool2d_with_indices (function) + * torch::nn::functional::fractional_max_pool3d_with_indices (function) + * torch::nn::functional::multi_head_attention_forward (function) + +torch::nn::init +--------------- + * torch::nn::init::calculate_gain (function) + * torch::nn::init::constant_ (function) + * torch::nn::init::dirac_ (function) + * torch::nn::init::eye_ (function) + * torch::nn::init::kaiming_normal_ (function) + * torch::nn::init::kaiming_uniform_ (function) + * torch::nn::init::normal_ (function) + * torch::nn::init::ones_ (function) + * torch::nn::init::orthogonal_ (function) + * torch::nn::init::sparse_ (function) + * torch::nn::init::uniform_ (function) + * torch::nn::init::xavier_normal_ (function) + * torch::nn::init::xavier_uniform_ (function) + * torch::nn::init::zeros_ (function) + +torch::nn::utils +---------------- + * torch::nn::utils::clip_grad_norm_ (function) + * torch::nn::utils::clip_grad_value_ (function) + * torch::nn::utils::parameters_to_vector (function) + * torch::nn::utils::rnn::invert_permutation (function) + * torch::nn::utils::vector_to_parameters (function) + +torch::optim +------------ + * torch::optim::AdagradOptions (struct) + * torch::optim::AdagradParamState (struct) + * torch::optim::AdamOptions (struct) + * torch::optim::AdamParamState (struct) + * torch::optim::AdamWOptions (struct) + * torch::optim::AdamWParamState (struct) + * torch::optim::LBFGSOptions (struct) + * torch::optim::LBFGSParamState (struct) + * torch::optim::RMSpropOptions (struct) + * torch::optim::RMSpropParamState (struct) + * torch::optim::SGDOptions (struct) + * torch::optim::SGDParamState (struct) + +torch::special +-------------- + * torch::special::airy_ai (function) + * torch::special::bessel_j0 (function) + * torch::special::bessel_j1 (function) + * torch::special::bessel_y0 (function) + * torch::special::bessel_y1 (function) + * torch::special::chebyshev_polynomial_t (function) + * torch::special::chebyshev_polynomial_u (function) + * torch::special::chebyshev_polynomial_v (function) + * torch::special::chebyshev_polynomial_w (function) + * torch::special::digamma (function) + * torch::special::entr (function) + * torch::special::erf (function) + * torch::special::erfc (function) + * torch::special::erfcx (function) + * torch::special::erfinv (function) + * torch::special::exp2 (function) + * torch::special::expit (function) + * torch::special::expm1 (function) + * torch::special::gammainc (function) + * torch::special::gammaincc (function) + * torch::special::gammaln (function) + * torch::special::hermite_polynomial_h (function) + * torch::special::hermite_polynomial_he (function) + * torch::special::i0 (function) + * torch::special::i0e (function) + * torch::special::i1 (function) + * torch::special::i1e (function) + * torch::special::laguerre_polynomial_l (function) + * torch::special::legendre_polynomial_p (function) + * torch::special::log1p (function) + * torch::special::log_ndtr (function) + * torch::special::log_softmax (function) + * torch::special::logit (function) + * torch::special::logsumexp (function) + * torch::special::modified_bessel_i0 (function) + * torch::special::modified_bessel_i1 (function) + * torch::special::modified_bessel_k0 (function) + * torch::special::modified_bessel_k1 (function) + * torch::special::multigammaln (function) + * torch::special::ndtr (function) + * torch::special::ndtri (function) + * torch::special::polygamma (function) + * torch::special::psi (function) + * torch::special::round (function) + * torch::special::scaled_modified_bessel_k0 (function) + * torch::special::scaled_modified_bessel_k1 (function) + * torch::special::shifted_chebyshev_polynomial_t (function) + * torch::special::shifted_chebyshev_polynomial_u (function) + * torch::special::shifted_chebyshev_polynomial_v (function) + * torch::special::shifted_chebyshev_polynomial_w (function) + * torch::special::sinc (function) + * torch::special::softmax (function) + * torch::special::spherical_bessel_j0 (function) + * torch::special::xlog1py (function) + * torch::special::xlogy (function) + * torch::special::zeta (function) + + +================================================== +Summary +================================================== + +Total APIs discovered: 699 +Documented: 472 +Missing: 227 +Coverage: 67.5% + +Category Covered Total % +---------------------------------------------------------------------- +at 4 4 100.0% +at::cuda 7 7 100.0% +c10 10 10 100.0% +c10::cuda 9 9 100.0% +c10::xpu 4 4 100.0% +other 1 1 100.0% +torch (core) 4 4 100.0% +torch (device) 5 16 31.2% +torch (macros) 3 5 60.0% +torch::autograd 3 3 100.0% +torch::data 29 34 85.3% +torch::fft 0 22 0.0% +torch::nn (modules) 243 319 76.2% +torch::nn::functional 100 124 80.6% +torch::nn::init 0 14 0.0% +torch::nn::utils 0 5 0.0% +torch::optim 13 25 52.0% +torch::serialize 4 4 100.0% +torch::special 0 56 0.0% +torch::stable 33 33 100.0% diff --git a/docs/cpp/cpp_html_issues.txt b/docs/cpp/cpp_html_issues.txt new file mode 100644 index 0000000000000..754a4bd2bc4c8 --- /dev/null +++ b/docs/cpp/cpp_html_issues.txt @@ -0,0 +1,4 @@ +HTML Formatting Check +================================================== + +No issues found. diff --git a/docs/cpp/source/Doxyfile b/docs/cpp/source/Doxyfile index 72e83733b6b4c..748f7a69ea266 100644 --- a/docs/cpp/source/Doxyfile +++ b/docs/cpp/source/Doxyfile @@ -34,6 +34,7 @@ INPUT = ../../../aten/src/ATen/ATen.h \ ../../../aten/src/ATen/core/ivalue.h \ ../../../aten/src/ATen/core/ScalarType.h \ ../../../aten/src/ATen/cuda/CUDAContext.h \ + ../../../aten/src/ATen/cuda/CUDAContextLight.h \ ../../../aten/src/ATen/cudnn/Descriptors.h \ ../../../aten/src/ATen/cudnn/Handles.h \ ../../../aten/src/ATen/cudnn/Types.h \ @@ -48,7 +49,9 @@ INPUT = ../../../aten/src/ATen/ATen.h \ ../../../build/aten/src/ATen/Functions.h \ ../../../build/aten/src/ATen/core/TensorBody.h \ ../../../c10/core/Device.h \ + ../../../c10/core/DeviceGuard.h \ ../../../c10/core/DeviceType.h \ + ../../../c10/core/Stream.h \ ../../../c10/util/Half.h \ ../../../c10/util/ArrayRef.h \ ../../../c10/util/OptionalArrayRef.h \ @@ -100,7 +103,8 @@ GENERATE_XML = YES # Set to NO if you do not want the Doxygen program listing included. # NOTE: setting to NO may result in unrecovered file relationships # (which file defined which compound) -XML_PROGRAMLISTING = YES +# Setting to NO to prevent inline code from polluting documentation +XML_PROGRAMLISTING = NO ################################################################################ # Doxygen preprocessor / parser control. # ################################################################################ @@ -129,8 +133,16 @@ EXTRACT_ALL = YES EXTRACT_PACKAGE = YES EXTRACT_STATIC = YES CASE_SENSE_NAMES = NO -EXCLUDE_SYMBOLS = caffe2::* cereal* DL* TH* cudnn* std::* +EXCLUDE_SYMBOLS = caffe2::* cereal* DL* TH* cudnn* std::* torch::nested* # EXCLUDE_SYMBOLS = c10::* caffe2::* cereal* DL* TH* cudnn* std::* + +# Don't include inline source code in documentation +INLINE_SOURCES = NO +# Don't include code bodies in detailed description +SOURCE_BROWSER = NO +# Don't reference source files +REFERENCES_LINK_SOURCE = NO +VERBATIM_HEADERS = NO ################################################################################ # Docstring control / customization. # ################################################################################ @@ -148,10 +160,21 @@ JAVADOC_AUTOBRIEF = YES # * # * \endrst # */ +# Or with /// style: +# /// \rst +# /// .. code-block:: cpp +# /// +# /// int main() { +# /// return 0; +# /// } +# /// +# /// \endrst +# # NOTE: # 1. \rst and \endrst must be on their own line. -# 2. leading-asterisk required. -ALIASES = "rst=\verbatim embed:rst:leading-asterisk" +# 2. For /// comments, use leading-slashes; for /** */ use leading-asterisk +# 3. PyTorch uses /// style comments primarily +ALIASES = "rst=\verbatim embed:rst:leading-slashes" ALIASES += "endrst=\endverbatim" ################################################################################ # Warning suppression. # diff --git a/docs/cpp/source/api/aten/accessors.md b/docs/cpp/source/api/aten/accessors.md new file mode 100644 index 0000000000000..83fefce724e7e --- /dev/null +++ b/docs/cpp/source/api/aten/accessors.md @@ -0,0 +1,49 @@ +--- +myst: + html_meta: + description: Tensor accessors in PyTorch C++ — efficient element-wise access to tensor data without overhead. + keywords: PyTorch, C++, tensor accessor, packed_accessor, data access +--- + +# Tensor Accessors + +For element-wise operations in custom kernels, use *accessors* to avoid +dynamic dispatch overhead. + +## CPU Accessors + +```cpp +torch::Tensor foo = torch::rand({12, 12}); + +// Create accessor - validates type and dimensions once +auto foo_a = foo.accessor(); + +float trace = 0; +for (int i = 0; i < foo_a.size(0); i++) { + trace += foo_a[i][i]; +} +``` + +## CUDA Packed Accessors + +For CUDA kernels, use *packed accessors* which copy metadata instead of +pointing to it: + +```cpp +__global__ void kernel(torch::PackedTensorAccessor64 foo, float* trace) { + int i = threadIdx.x; + gpuAtomicAdd(trace, foo[i][i]); +} + +torch::Tensor foo = torch::rand({12, 12}).cuda(); +auto foo_a = foo.packed_accessor64(); + +float trace = 0; +kernel<<<1, 12>>>(foo_a, &trace); +``` + +```{tip} + +Use `PackedTensorAccessor32` and `packed_accessor32` for 32-bit indexing, +which is faster on CUDA but may overflow for large tensors. +``` diff --git a/docs/cpp/source/api/aten/creation.md b/docs/cpp/source/api/aten/creation.md new file mode 100644 index 0000000000000..27ddcd616db55 --- /dev/null +++ b/docs/cpp/source/api/aten/creation.md @@ -0,0 +1,235 @@ +--- +myst: + html_meta: + description: Tensor creation functions in PyTorch C++ — zeros, ones, randn, arange, from_blob, and more. + keywords: PyTorch, C++, tensor creation, zeros, ones, randn, arange, from_blob +--- + +# Tensor Creation + +Factory functions create new tensors with different initialization patterns. +All factory functions follow a general schema: + +```cpp +torch::(, , ) +``` + +## Available Factory Functions + +- `torch::zeros` - Tensor filled with zeros +- `torch::ones` - Tensor filled with ones +- `torch::empty` - Uninitialized tensor +- `torch::full` - Tensor filled with a single value +- `torch::rand` - Uniform random values on [0, 1) +- `torch::randn` - Standard normal distribution +- `torch::randint` - Random integers in a range +- `torch::arange` - Sequence of integers +- `torch::linspace` - Linearly spaced values +- `torch::logspace` - Logarithmically spaced values +- `torch::eye` - Identity matrix +- `torch::randperm` - Random permutation of integers + +## Specifying a Size + +Functions that do not require specific arguments can be invoked with just a +size. For example, the following line creates a vector with 5 components: + +```cpp +torch::Tensor tensor = torch::ones(5); +``` + +An `IntArrayRef` is constructed by specifying the size along each dimension in +curly braces. For example, `{2, 3}` for a matrix with two rows and three +columns, `{3, 4, 5}` for a three-dimensional tensor: + +```cpp +torch::Tensor tensor = torch::randn({3, 4, 5}); +assert(tensor.sizes() == std::vector{3, 4, 5}); +``` + +You can also pass an `std::vector` instead of curly braces. +Use `tensor.size(i)` to access a single dimension. + +## Passing Function-Specific Parameters + +Some factory functions accept additional parameters. For example, `randint` +takes an upper bound on the value for the integers it generates: + +```cpp +torch::Tensor tensor = torch::randint(/*high=*/10, {5, 5}); + +// With a lower bound +torch::Tensor tensor = torch::randint(/*low=*/3, /*high=*/10, {5, 5}); +``` + +```{tip} + +The size always follows the function-specific arguments. +``` + +```{attention} + +Some functions like `arange` do not need a size at all, since it is fully +determined by the function-specific arguments (the range bounds). +``` + +## Configuring Properties with TensorOptions + +`TensorOptions` configures the data type, layout, device, and +`requires_grad` of a new tensor. The construction axes are: + +- `dtype`: the data type of the elements (e.g. `kFloat32`, `kInt64`) +- `layout`: either `kStrided` (dense) or `kSparse` +- `device`: a compute device (e.g. `kCPU`, `kCUDA`) +- `requires_grad`: whether to track gradients + +Allowed values: + +- `dtype`: `kUInt8`, `kInt8`, `kInt16`, `kInt32`, `kInt64`, + `kFloat32`, `kFloat64` +- `layout`: `kStrided`, `kSparse` +- `device`: `kCPU`, or `kCUDA` (with an optional device index) +- `requires_grad`: `true` or `false` + +```{tip} + +Rust-style shorthands exist for dtypes, like `kF32` instead of +`kFloat32`. See +[torch/types.h](https://github.com/pytorch/pytorch/blob/main/torch/csrc/api/include/torch/types.h) +for the full list. +``` + +Here is an example of creating a `TensorOptions` object: + +```cpp +auto options = + torch::TensorOptions() + .dtype(torch::kFloat32) + .layout(torch::kStrided) + .device(torch::kCUDA, 1) + .requires_grad(true); + +torch::Tensor tensor = torch::full({3, 4}, /*value=*/123, options); + +assert(tensor.dtype() == torch::kFloat32); +assert(tensor.layout() == torch::kStrided); +assert(tensor.device().type() == torch::kCUDA); +assert(tensor.device().index() == 1); +assert(tensor.requires_grad()); +``` + +**Defaults:** Any axis you omit takes its default value: `kFloat32` for dtype, +`kStrided` for layout, `kCPU` for device, and `false` for +`requires_grad`. This means you can omit `TensorOptions` entirely: + +```cpp +// A 32-bit float, strided, CPU tensor that does not require a gradient. +torch::Tensor tensor = torch::randn({3, 4}); +``` + +**Shorthand syntax:** For each axis there is a free function in the `torch::` +namespace (`torch::dtype()`, `torch::device()`, `torch::layout()`, +`torch::requires_grad()`). Each returns a `TensorOptions` object that can +be further refined with builder methods: + +```cpp +// These are equivalent: +torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32)) +torch::ones(10, torch::dtype(torch::kFloat32)) + +// Chaining: +torch::ones(10, torch::dtype(torch::kFloat32).layout(torch::kStrided)) +``` + +**Implicit construction:** `TensorOptions` is implicitly constructible from +individual values, so when only one axis differs from the default you can +write: + +```cpp +torch::ones(10, torch::kFloat32) +``` + +Putting it all together, a C++ tensor creation call mirrors the Python +equivalent closely: + +```python +# Python +torch.randn(3, 4, dtype=torch.float32, device=torch.device('cuda', 1), requires_grad=True) +``` + +```cpp +// C++ +torch::randn({3, 4}, torch::dtype(torch::kFloat32).device(torch::kCUDA, 1).requires_grad(true)) +``` + +## Using Externally Created Data + +If you already have tensor data allocated in memory (CPU or CUDA), use +`from_blob` to view that memory as a `Tensor`: + +```cpp +float data[] = {1, 2, 3, 4, 5, 6}; +torch::Tensor tensor = torch::from_blob(data, {2, 3}); +``` + +```{note} + +Tensors created with `from_blob` cannot be resized because ATen does not +own the memory. +``` + +## Tensor Conversion + +Use `to()` to convert tensors between dtypes and devices. The conversion +creates a new tensor and does not occur in-place: + +```cpp +torch::Tensor source = torch::randn({2, 3}, torch::kInt64); + +// Convert dtype +torch::Tensor float_tensor = source.to(torch::kFloat32); + +// Move to GPU (default CUDA device) +torch::Tensor gpu_tensor = float_tensor.to(torch::kCUDA); + +// Specific GPU device +torch::Tensor gpu1_tensor = float_tensor.to(torch::Device(torch::kCUDA, 1)); + +// Async copy +torch::Tensor async_tensor = gpu_tensor.to(torch::kCPU, /*non_blocking=*/true); +``` + +```{attention} + +The result of the conversion is a new tensor pointing to new memory, +unrelated to the source tensor. +``` + +## Scalars and Zero-Dimensional Tensors + +`Scalar` represents a single dynamically-typed number. Like a `Tensor`, +`Scalar` is dynamically typed and can hold any of ATen's number types. +Scalars can be implicitly constructed from C++ number types: + +```cpp +namespace torch { +Tensor addmm(Scalar beta, const Tensor & self, + Scalar alpha, const Tensor & mat1, + const Tensor & mat2); +Scalar sum(const Tensor & self); +} // namespace torch + +// Usage +torch::Tensor a = ...; +torch::Tensor b = ...; +torch::Tensor c = ...; +torch::Tensor r = torch::addmm(1.0, a, .5, b, c); +``` + +Zero-dimensional tensors hold a single value and can reference elements in +larger tensors: + +```cpp +torch::Tensor matrix = torch::rand({10, 20}); +matrix[1][2] = 4; // matrix[1][2] is a zero-dimensional tensor +``` diff --git a/docs/cpp/source/api/aten/index.md b/docs/cpp/source/api/aten/index.md new file mode 100644 index 0000000000000..e9980c4cca39f --- /dev/null +++ b/docs/cpp/source/api/aten/index.md @@ -0,0 +1,76 @@ +--- +myst: + html_meta: + description: ATen C++ API — the foundational tensor library powering PyTorch, including tensor operations, types, and CUDA support. + keywords: PyTorch, C++, ATen, tensor, operations, CUDA +--- + +# ATen: Tensor Library + +ATen (A Tensor Library) is the foundational tensor and mathematical operation +library on which all of PyTorch is built. It provides the core `Tensor` class +and hundreds of mathematical operations that work on tensors. + +**When to use ATen directly:** + +- When writing low-level tensor operations or custom kernels +- When you need direct access to tensor data and metadata +- When working with the PyTorch internals or extending PyTorch + +**Basic usage:** + +```cpp +#include + +// Create tensors +at::Tensor a = at::ones({2, 3}); +at::Tensor b = at::randn({2, 3}); + +// Perform operations +at::Tensor c = a + b; +at::Tensor d = at::matmul(a.t(), b); + +// Move to GPU +if (at::cuda::is_available()) { + at::Tensor gpu_tensor = c.to(at::kCUDA); +} +``` + +For most applications, prefer using the higher-level `torch::` namespace +(see {doc}`../nn/index`, {doc}`../optim/index`) which provides a more user-friendly API. + +## Header Files + +The following headers are part of the ATen public API: + +- `ATen/ATen.h` - Main ATen header +- `ATen/Backend.h` - Backend enumeration +- `ATen/core/Tensor.h` - Tensor class +- `ATen/core/ivalue.h` - IValue type (see {doc}`../c10/index`) +- `ATen/core/ScalarType.h` - Data type definitions +- `ATen/TensorOptions.h` - Tensor creation options +- `ATen/Scalar.h` - Scalar type +- `ATen/Layout.h` - Tensor layout +- `ATen/DeviceGuard.h` - Device context management +- `ATen/native/TensorShape.h` - Tensor shape operations +- `ATen/cuda/CUDAContext.h` - CUDA context (see {doc}`../cuda/index`) +- `ATen/cudnn/Descriptors.h` - cuDNN descriptors +- `ATen/mkl/Descriptors.h` - MKL descriptors + +```{note} + +The core `at::Tensor` class is defined in a generated header file +(`TensorBody.h`) that only exists after building PyTorch. The documentation +below describes the API manually. +``` + +## ATen Categories + +```{toctree} +:maxdepth: 1 + +tensor +creation +indexing +accessors +``` diff --git a/docs/cpp/source/api/aten/indexing.md b/docs/cpp/source/api/aten/indexing.md new file mode 100644 index 0000000000000..63481cb1cd5a4 --- /dev/null +++ b/docs/cpp/source/api/aten/indexing.md @@ -0,0 +1,193 @@ +--- +myst: + html_meta: + description: Tensor indexing in PyTorch C++ — Slice, None, Ellipsis, boolean masks, and advanced indexing with index() and index_put_(). + keywords: PyTorch, C++, tensor indexing, Slice, index, index_put_, boolean mask, fancy indexing +--- + +# Tensor Indexing + +The PyTorch C++ API provides tensor indexing similar to Python. Use +`torch::indexing` namespace for index types: + +```cpp +using namespace torch::indexing; +``` + +The main difference from Python is that instead of using the `[]` operator, +the C++ API uses the `index` and `index_put_` methods: + +- `torch::Tensor::index` — read elements +- `torch::Tensor::index_put_` — write elements + +## Index Types + +The `TensorIndex` class accepts six types of indices via implicit constructors: + +```{list-table} +:widths: 25 35 40 +:header-rows: 1 + +* - Type + - C++ + - Python equivalent +* - None (unsqueeze) + - `None` + - `None` +* - Ellipsis + - `Ellipsis` or `"..."` + - `...` +* - Integer + - `0`, `1`, `-1` + - `0`, `1`, `-1` +* - Boolean + - `true`, `false` + - `True`, `False` +* - Slice + - `Slice(start, stop, step)` + - `start:stop:step` +* - Tensor + - `torch::tensor({0, 2})` + - `torch.tensor([0, 2])` +``` + +## Getter Operations + +```{list-table} +:widths: 40 60 +:header-rows: 1 + +* - Python + - C++ +* - `tensor[0]` + - `tensor.index({0})` +* - `tensor[-1]` + - `tensor.index({-1})` +* - `tensor[1, 2]` + - `tensor.index({1, 2})` +* - `tensor[1, :, 3]` + - `tensor.index({1, Slice(), 3})` +* - `tensor[None]` + - `tensor.index({None})` +* - `tensor[:, None]` + - `tensor.index({Slice(), None})` +* - `tensor[...]` + - `tensor.index({Ellipsis})` or `tensor.index({"..."})` +* - `tensor[..., 0]` + - `tensor.index({Ellipsis, 0})` +* - `tensor[1::2]` + - `tensor.index({Slice(1, None, 2)})` +* - `tensor[True]` + - `tensor.index({true})` +* - `tensor[torch.tensor([1, 2])]` + - `tensor.index({torch::tensor({1, 2})})` +* - `tensor[bool_mask]` + - `tensor.index({bool_mask})` +* - `tensor[:, torch.tensor([[0,1],[4,3]])]` + - `tensor.index({Slice(), torch::tensor({{0,1},{4,3}})})` +* - `tensor[cond > 0]` + - `tensor.index({cond > 0})` +``` + +## Setter Operations + +```{list-table} +:widths: 40 60 +:header-rows: 1 + +* - Python + - C++ +* - `tensor[0] = 1` + - `tensor.index_put_({0}, 1)` +* - `tensor[1, 2] = 1` + - `tensor.index_put_({1, 2}, 1)` +* - `tensor[1] = torch.arange(5)` + - `tensor.index_put_({1}, torch::arange(5))` +* - `tensor[1::2] = 1` + - `tensor.index_put_({Slice(1, None, 2)}, 1)` +* - `tensor[0, 1::2] = torch.tensor([3., 4.])` + - `tensor.index_put_({0, Slice(1, None, 2)}, torch::tensor({3., 4.}))` +* - `tensor[...] = 0` + - `tensor.index_put_({Ellipsis}, 0)` +* - `tensor[None] = value` + - `tensor.index_put_({None}, value)` +* - `tensor[bool_mask] = 0` + - `tensor.index_put_({bool_mask}, 0)` +* - `tensor[torch.tensor([0, 2])] = value` + - `tensor.index_put_({torch::tensor({0, 2})}, value)` +* - `tensor[1:2, torch.tensor([1,2])] = 0` + - `tensor.index_put_({Slice(1, 2), torch::tensor({1, 2})}, 0)` +``` + +The `index_put_` method also accepts an optional `accumulate` parameter. +When `true`, values are added to existing values instead of replacing them: + +```cpp +tensor.index_put_({mask}, values, /*accumulate=*/true); +``` + +## Slice Syntax + +The `Slice` constructor signature is: + +```cpp +Slice( + std::optional start = std::nullopt, + std::optional stop = std::nullopt, + std::optional step = std::nullopt); +``` + +Pass `None` for open-ended bounds: + +```{list-table} +:widths: 30 70 +:header-rows: 1 + +* - Python + - C++ +* - `:` or `::` + - `Slice()` or `Slice(None, None)` +* - `1:` + - `Slice(1, None)` +* - `:3` + - `Slice(None, 3)` +* - `1:3` + - `Slice(1, 3)` +* - `1:3:2` + - `Slice(1, 3, 2)` +* - `::2` + - `Slice(None, None, 2)` +``` + +## Full Example + +```cpp +#include + +using namespace torch::indexing; + +auto tensor = torch::arange(2 * 3 * 4).reshape({2, 3, 4}); + +// Basic indexing +auto row = tensor.index({0}); // tensor[0] +auto elem = tensor.index({1, 2, 3}); // tensor[1, 2, 3] + +// Slicing +auto sliced = tensor.index({Slice(), Slice(0, 2)}); // tensor[:, 0:2] + +// None (unsqueeze) and Ellipsis +auto unsqueezed = tensor.index({None}); // tensor[None] +auto last_dim = tensor.index({Ellipsis, -1}); // tensor[..., -1] + +// Boolean mask indexing +auto mask = tensor > 10; +auto selected = tensor.index({mask}); // tensor[tensor > 10] + +// Integer tensor (fancy) indexing +auto idx = torch::tensor({0, 2}); +auto gathered = tensor.index({Slice(), idx}); // tensor[:, [0, 2]] + +// Setting values +tensor.index_put_({0, Slice(), 0}, 99); // tensor[0, :, 0] = 99 +tensor.index_put_({mask}, 0); // tensor[tensor > 10] = 0 +``` diff --git a/docs/cpp/source/api/aten/tensor.md b/docs/cpp/source/api/aten/tensor.md new file mode 100644 index 0000000000000..227c94e13d7ff --- /dev/null +++ b/docs/cpp/source/api/aten/tensor.md @@ -0,0 +1,288 @@ +--- +myst: + html_meta: + description: at::Tensor class reference — the primary tensor type in PyTorch C++ with creation, indexing, device, and dtype APIs. + keywords: PyTorch, C++, Tensor, at::Tensor, TensorOptions, ScalarType, DeviceGuard +--- + +# Tensor Class + +The `at::Tensor` class is the primary tensor class in ATen, representing +a multi-dimensional array with a specific data type and device. + +## Tensor + +```{cpp:class} at::Tensor + +The primary tensor class in ATen. Represents a multi-dimensional array +with a specific data type and device. +``` + +```{cpp:function} Tensor() + +Default constructor. Creates an undefined tensor. +``` + +```{cpp:function} int64_t dim() const + +Returns the number of dimensions of the tensor. +``` + +```{cpp:function} int64_t size(int64_t dim) const + +Returns the size of the tensor at the given dimension. +``` + +```{cpp:function} IntArrayRef sizes() const + +Returns the sizes of all dimensions. +``` + +```{cpp:function} IntArrayRef strides() const + +Returns the strides of all dimensions. +``` + +```{cpp:function} ScalarType scalar_type() const + +Returns the data type of the tensor. +``` + +```{cpp:function} Device device() const + +Returns the device where the tensor is stored. +``` + +```{cpp:function} bool is_cuda() const + +Returns true if the tensor is on a CUDA device. +``` + +```{cpp:function} bool is_cpu() const + +Returns true if the tensor is on CPU. +``` + +```{cpp:function} bool requires_grad() const + +Returns true if gradients need to be computed for this tensor. +``` + +```{cpp:function} Tensor& requires_grad_(bool requires_grad = true) + +Sets whether gradients should be computed for this tensor. +``` + +```{cpp:function} Tensor to(Device device) const + +Returns a tensor on the specified device. +``` + +```{cpp:function} Tensor to(ScalarType dtype) const + +Returns a tensor with the specified data type. +``` + +```{cpp:function} Tensor contiguous() const + +Returns a contiguous tensor with the same data. +``` + +```{cpp:function} void* data_ptr() const + +Returns a pointer to the underlying data. +``` + +**Example:** + +```cpp +#include + +at::Tensor a = at::ones({2, 2}, at::kInt); +at::Tensor b = at::randn({2, 2}); +auto c = a + b.to(at::kInt); +``` + +## TensorOptions + +```{cpp:class} at::TensorOptions + +A class to specify options for tensor creation, including dtype, device, +layout, and requires_grad. +``` + +```{cpp:function} TensorOptions() + +Default constructor. +``` + +```{cpp:function} TensorOptions dtype(ScalarType dtype) const + +Returns options with the specified data type. +``` + +```{cpp:function} TensorOptions device(Device device) const + +Returns options with the specified device. +``` + +```{cpp:function} TensorOptions layout(Layout layout) const + +Returns options with the specified layout. +``` + +```{cpp:function} TensorOptions requires_grad(bool requires_grad) const + +Returns options with the specified requires_grad setting. +``` + +**Example:** + +```cpp +auto options = at::TensorOptions() + .dtype(at::kFloat) + .device(at::kCUDA, 0) + .requires_grad(true); + +at::Tensor tensor = at::zeros({3, 4}, options); +``` + +## Scalar + +```{cpp:class} at::Scalar + +Represents a scalar value that can be converted to various numeric types. +``` + +```{cpp:function} Scalar(int64_t v) + +Construct from an integer. +``` + +```{cpp:function} Scalar(double v) + +Construct from a double. +``` + +```{cpp:function} template T to() const + +Convert to the specified type. +``` + +```{cpp:function} bool isIntegral(bool includeBool = false) const + +Returns true if the scalar is an integral type. +``` + +```{cpp:function} bool isFloatingPoint() const + +Returns true if the scalar is a floating point type. +``` + +## ScalarType + +```{cpp:enum-class} at::ScalarType + +Enumeration of data types supported by tensors. +``` + +```{cpp:enumerator} Byte + +8-bit unsigned integer (uint8_t) +``` + +```{cpp:enumerator} Char + +8-bit signed integer (int8_t) +``` + +```{cpp:enumerator} Short + +16-bit signed integer (int16_t) +``` + +```{cpp:enumerator} Int + +32-bit signed integer (int32_t) +``` + +```{cpp:enumerator} Long + +64-bit signed integer (int64_t) +``` + +```{cpp:enumerator} Half + +16-bit floating point (float16) +``` + +```{cpp:enumerator} Float + +32-bit floating point (float) +``` + +```{cpp:enumerator} Double + +64-bit floating point (double) +``` + +```{cpp:enumerator} Bool + +Boolean type +``` + +```{cpp:enumerator} BFloat16 + +Brain floating point (bfloat16) +``` + +Convenience constants: + +- `at::kByte`, `at::kChar`, `at::kShort`, `at::kInt`, `at::kLong` +- `at::kHalf`, `at::kFloat`, `at::kDouble`, `at::kBFloat16` +- `at::kBool` + +## DeviceGuard + +```{doxygenclass} c10::DeviceGuard +:members: +:undoc-members: +``` + +**Example:** + +```cpp +{ + c10::DeviceGuard guard(at::Device(at::kCUDA, 1)); + // Operations here run on CUDA device 1 + auto tensor = at::zeros({2, 2}); +} +// Previous device is restored +``` + +## Layout + +```{cpp:enum-class} at::Layout + +Enumeration of tensor memory layouts. +``` + +```{cpp:enumerator} Strided + +Dense tensor with strides. +``` + +```{cpp:enumerator} Sparse + +Sparse tensor (COO format). +``` + +```{cpp:enumerator} SparseCsr + +Sparse tensor in CSR format. +``` + +```{cpp:enumerator} SparseCsc + +Sparse tensor in CSC format. +``` diff --git a/docs/cpp/source/api/autograd/custom_functions.md b/docs/cpp/source/api/autograd/custom_functions.md new file mode 100644 index 0000000000000..475c975cdd9aa --- /dev/null +++ b/docs/cpp/source/api/autograd/custom_functions.md @@ -0,0 +1,95 @@ +--- +myst: + html_meta: + description: Custom autograd functions in PyTorch C++ — defining forward and backward passes with torch::autograd::Function. + keywords: PyTorch, C++, autograd, custom function, forward, backward, Function +--- + +# Custom Autograd Functions + +PyTorch allows you to define custom autograd functions with custom forward +and backward implementations. + +## Function Base Class + +```{doxygenstruct} torch::autograd::Function +:members: +:undoc-members: +``` + +## AutogradContext + +```{doxygenstruct} torch::autograd::AutogradContext +:members: +:undoc-members: +``` + +## Creating Custom Functions + +To create a custom autograd function, inherit from `torch::autograd::Function` +and implement the static `forward` and `backward` methods: + +**Example:** + +```cpp +class MyReLU : public torch::autograd::Function { + public: + static torch::Tensor forward( + torch::autograd::AutogradContext* ctx, + torch::Tensor input) { + ctx->save_for_backward({input}); + return input.clamp_min(0); + } + + static torch::autograd::variable_list backward( + torch::autograd::AutogradContext* ctx, + torch::autograd::variable_list grad_outputs) { + auto saved = ctx->get_saved_variables(); + auto input = saved[0]; + auto grad_output = grad_outputs[0]; + auto grad_input = grad_output * (input > 0).to(grad_output.dtype()); + return {grad_input}; + } +}; + +// Usage +auto output = MyReLU::apply(input); +``` + +## Custom Kernels and AutoDispatchBelowADInplaceOrView + +For users implementing custom kernels who want to redispatch below `Autograd` dispatch +keys, use `at::AutoDispatchBelowADInplaceOrView` instead of `InferenceMode`: + +```cpp +class ROIAlignFunction : public torch::autograd::Function { + public: + static torch::autograd::variable_list forward( + torch::autograd::AutogradContext* ctx, + const torch::autograd::Variable& input, + const torch::autograd::Variable& rois, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width, + int64_t sampling_ratio, + bool aligned) { + ctx->saved_data["spatial_scale"] = spatial_scale; + ctx->saved_data["pooled_height"] = pooled_height; + ctx->saved_data["pooled_width"] = pooled_width; + ctx->saved_data["sampling_ratio"] = sampling_ratio; + ctx->saved_data["aligned"] = aligned; + ctx->saved_data["input_shape"] = input.sizes(); + ctx->save_for_backward({rois}); + + at::AutoDispatchBelowADInplaceOrView guard; + auto result = roi_align( + input, rois, spatial_scale, pooled_height, + pooled_width, sampling_ratio, aligned); + return {result}; + } +}; +``` + +For customized inplace and view kernels, see the +[custom kernel tutorial](https://pytorch.org/tutorials/advanced/cpp_extension.html#backward-pass) +for more details. diff --git a/docs/cpp/source/api/autograd/gradient.md b/docs/cpp/source/api/autograd/gradient.md new file mode 100644 index 0000000000000..6f595496fb7a0 --- /dev/null +++ b/docs/cpp/source/api/autograd/gradient.md @@ -0,0 +1,86 @@ +--- +myst: + html_meta: + description: PyTorch C++ API for computing gradients — torch::autograd::backward and torch::autograd::grad functions for automatic differentiation. + keywords: PyTorch, C++, autograd, backward, grad, gradient, automatic differentiation +--- + +# Gradient Computation + +PyTorch provides functions for computing gradients of tensors with respect +to graph leaves. + +## Gradient Functions + +```{cpp:function} void torch::autograd::backward(const variable_list& tensors, const variable_list& grad_tensors = {}, std::optional retain_graph = std::nullopt, bool create_graph = false, const variable_list& inputs = {}) + +Computes the sum of gradients of given tensors with respect to graph leaves. + +The graph is differentiated using the chain rule. If any of `tensors` +are non-scalar (i.e. their data has more than one element) and require +gradient, then the Jacobian-vector product would be computed, in this case +the function additionally requires specifying `grad_tensors`. It should be a +sequence of matching length, that contains the "vector" in the +Jacobian-vector product, usually the gradient of the differentiated function +w.r.t. corresponding tensors (`torch::Tensor()` is an acceptable value for +all tensors that don't need gradient tensors). + +This function accumulates gradients in the leaves — you might need to zero +them before calling it. + +:param tensors: Tensors of which the derivative will be computed. +:param grad_tensors: The "vector" in the Jacobian-vector product, usually + gradients w.r.t. each element of corresponding tensors. + `torch::Tensor()` values can be specified for scalar Tensors or ones + that don't require grad. If a `torch::Tensor()` value would be + acceptable for all grad_tensors, then this argument is optional. +:param retain_graph: If `false`, the graph used to compute the grad will + be freed. Note that in nearly all cases setting this option to `true` + is not needed and often can be worked around in a much more efficient + way. Defaults to the value of `create_graph`. +:param create_graph: If `true`, graph of the derivative will be + constructed, allowing to compute higher order derivative products. + Defaults to `false`. +:param inputs: Inputs w.r.t. which the gradient will be accumulated into + `at::Tensor::grad`. All other Tensors will be ignored. If not + provided, the gradient is accumulated into all the leaf Tensors that + were used to compute `tensors`. +``` + +```{doxygenfunction} torch::autograd::grad +``` + +**Example:** + +```cpp +#include + +auto x = torch::randn({2, 2}, torch::requires_grad()); +auto y = x * x; +auto z = y.sum(); + +// Compute gradients +z.backward(); +std::cout << x.grad() << std::endl; + +// Or use grad() for specific outputs +auto grads = torch::autograd::grad({z}, {x}); +``` + +## Tensor Gradient Methods + +Tensors have built-in methods for gradient computation: + +```cpp +// Enable gradient tracking +auto x = torch::randn({2, 2}).requires_grad_(true); + +// Check if gradient is required +bool needs_grad = x.requires_grad(); + +// Access the gradient after backward +auto grad = x.grad(); + +// Detach from computation graph +auto x_detached = x.detach(); +``` diff --git a/docs/cpp/source/api/autograd/index.md b/docs/cpp/source/api/autograd/index.md new file mode 100644 index 0000000000000..37b3cb79f747e --- /dev/null +++ b/docs/cpp/source/api/autograd/index.md @@ -0,0 +1,56 @@ +--- +myst: + html_meta: + description: PyTorch C++ Autograd API — automatic differentiation for tensor computations. + keywords: PyTorch, C++, autograd, automatic differentiation, gradient +--- + +# Autograd: Automatic Differentiation + +PyTorch's autograd system provides automatic differentiation for all operations +on tensors. It records operations on tensors to build a computational graph, +then computes gradients automatically via backpropagation. + +**When to use Autograd:** + +- When training neural networks (gradients are computed automatically) +- When implementing custom backward passes for specialized operations +- When you need fine-grained control over gradient computation + +**Basic usage:** + +```cpp +#include + +// Create tensor with gradient tracking +auto x = torch::randn({2, 2}, torch::requires_grad()); +auto y = x * x; +auto z = y.sum(); + +// Compute gradients via backpropagation +z.backward(); +std::cout << x.grad() << std::endl; // dz/dx = 2x + +// Disable gradient tracking for inference +{ + torch::NoGradGuard no_grad; + auto result = model->forward(input); // No gradients computed +} +``` + +## Header Files + +- `torch/csrc/autograd/autograd.h` - High-level autograd API +- `torch/csrc/autograd/function.h` - Custom autograd functions +- `torch/csrc/autograd/grad_mode.h` - Gradient computation modes +- `torch/csrc/api/include/torch/autograd.h` - C++ Frontend autograd + +## Autograd Categories + +```{toctree} +:maxdepth: 1 + +gradient +custom_functions +modes +``` diff --git a/docs/cpp/source/api/autograd/modes.md b/docs/cpp/source/api/autograd/modes.md new file mode 100644 index 0000000000000..f1a281deae4f8 --- /dev/null +++ b/docs/cpp/source/api/autograd/modes.md @@ -0,0 +1,166 @@ +--- +myst: + html_meta: + description: Gradient mode guards in PyTorch C++ — NoGradGuard and InferenceMode for disabling gradient computation. + keywords: PyTorch, C++, NoGradGuard, InferenceMode, no_grad, inference, RAII guard +--- + +# Gradient Modes + +PyTorch provides RAII guards to control gradient computation behavior. + +## NoGradGuard + +```{cpp:class} torch::NoGradGuard + +RAII guard that disables gradient computation within its scope. +``` + +```{cpp:function} NoGradGuard() + +Disables gradient computation. +``` + +```{cpp:function} ~NoGradGuard() + +Restores previous gradient mode. +``` + +**Example:** + +```cpp +{ + torch::NoGradGuard no_grad; + // No gradients computed in this scope + auto result = model->forward(input); +} +``` + +## InferenceMode + +`c10::InferenceMode` is a RAII guard analogous to `NoGradMode` designed for use +when you are certain your operations will have no interactions with autograd +(e.g., model inference). Compared to `NoGradMode`, code run under this mode gets +better performance by disabling autograd-related work like view tracking and version +counter bumps. However, tensors created inside `InferenceMode` have more limitations +when interacting with the autograd system. + +```{cpp:class} c10::InferenceMode + +RAII guard that enables inference mode for optimized inference. +This is more efficient than NoGradGuard for inference-only workloads. +``` + +```{cpp:function} explicit InferenceMode(bool enabled = true) + +Enables or disables inference mode. +``` + +**Inference Tensors:** + +`InferenceMode` can be enabled for a given block of code. Inside `InferenceMode`, +all newly allocated (non-view) tensors are marked as inference tensors. Inference tensors: + +- Do not have a version counter, so an error will be raised if you try to read their version + (e.g., because you saved this tensor for backward). +- Are immutable outside `InferenceMode`. An error will be raised if you try to: + + - Mutate their data outside InferenceMode. + - Mutate them to `requires_grad=True` outside InferenceMode. + - To work around this, make a clone outside `InferenceMode` to get a normal tensor before mutating. + +A non-view tensor is an inference tensor if and only if it was allocated inside `InferenceMode`. +A view tensor is an inference tensor if and only if it is a view of an inference tensor. + +**Performance Guarantees:** + +Inside an `InferenceMode` block: + +- Like `NoGradMode`, all operations do not record `grad_fn` even if their inputs have + `requires_grad=True`. This applies to both inference tensors and normal tensors. +- View operations on inference tensors do not perform view tracking. View and non-view + inference tensors are indistinguishable. +- Inplace operations on inference tensors are guaranteed not to do a version bump. + +For more implementation details, see the [RFC-0011-InferenceMode](https://github.com/pytorch/rfcs/pull/17). + +**Basic Example:** + +```cpp +{ + c10::InferenceMode guard; + // Optimized inference without gradient tracking + auto result = model->forward(input); +} +``` + +**Inference Workload Example:** + +```cpp +c10::InferenceMode guard; +model.load_jit(saved_model); +auto inputs = preprocess_tensors(data); +auto out = model.forward(inputs); +auto outputs = postprocess_tensors(out); +``` + +**Nested InferenceMode:** + +Unlike some other guards, `InferenceMode` can be nested with different enabled/disabled states: + +```cpp +{ + c10::InferenceMode guard(true); + // InferenceMode is on + { + c10::InferenceMode guard(false); + // InferenceMode is off + } + // InferenceMode is on +} +// InferenceMode is off +``` + +## InferenceMode vs NoGradMode + +`InferenceMode` is preferred over `NoGradMode` for pure inference workloads because +it provides better performance. Key differences: + +- Both guards affect tensor execution to skip work not related to inference, but + `InferenceMode` also affects tensor creation while `NoGradMode` doesn't. +- Tensors created inside `InferenceMode` are marked as inference tensors with + certain limitations that apply after exiting `InferenceMode`. +- `InferenceMode` can be nested with enabled/disabled states. + +## Migrating from AutoNonVariableTypeMode + +The legacy `AutoNonVariableTypeMode` guard (now renamed to +`AutoDispatchBelowADInplaceOrView`) was commonly used for inference workloads +but is unsafe — it can silently bypass safety checks and produce wrong results. + +- **For inference-only workloads** (e.g. loading a pretrained JIT model and + running inference in C++ runtime), use `c10::InferenceMode` as a drop-in + replacement. It preserves the performance characteristics while providing + correctness guarantees. + +- **For custom autograd kernels** that need to redispatch below the Autograd + dispatch key, use `AutoDispatchBelowADInplaceOrView` instead: + + ```cpp + class ROIAlignFunction : public torch::autograd::Function { + public: + static torch::autograd::variable_list forward( + torch::autograd::AutogradContext* ctx, + const torch::autograd::Variable& input, + const torch::autograd::Variable& rois, + double spatial_scale, int64_t pooled_height, + int64_t pooled_width, int64_t sampling_ratio, bool aligned) { + ctx->saved_data["spatial_scale"] = spatial_scale; + ctx->save_for_backward({rois}); + at::AutoDispatchBelowADInplaceOrView guard; + auto result = roi_align(input, rois, spatial_scale, + pooled_height, pooled_width, sampling_ratio, aligned); + return {result}; + } + }; + ``` diff --git a/docs/cpp/source/api/c10/device.md b/docs/cpp/source/api/c10/device.md new file mode 100644 index 0000000000000..229082935c11b --- /dev/null +++ b/docs/cpp/source/api/c10/device.md @@ -0,0 +1,114 @@ +--- +myst: + html_meta: + description: Device and DeviceType in PyTorch C++ — c10::Device for specifying CPU, CUDA, XPU, and other backends. + keywords: PyTorch, C++, Device, DeviceType, CPU, CUDA, XPU, MPS, c10 +--- + +# Device and DeviceType + +PyTorch provides device abstractions for writing code that works across +CPU, CUDA, and other backends. + +## Device + +```{doxygenstruct} c10::Device +:members: +:undoc-members: +``` + +**Example:** + +```cpp +c10::Device cpu_device(c10::kCPU); +c10::Device cuda_device(c10::kCUDA, 0); // CUDA device 0 + +if (cuda_device.is_cuda()) { + std::cout << "Using CUDA device " << cuda_device.index() << std::endl; +} +``` + +## DeviceType + +```{cpp:enum-class} c10::DeviceType + +Enumeration of supported device types. +``` + +```{cpp:enumerator} CPU = 0 + +CPU device. +``` + +```{cpp:enumerator} CUDA = 1 + +NVIDIA CUDA GPU. +``` + +```{cpp:enumerator} HIP = 6 + +AMD HIP GPU. +``` + +```{cpp:enumerator} XLA = 9 + +XLA / TPU. +``` + +```{cpp:enumerator} Vulkan = 10 + +Vulkan GPU. +``` + +```{cpp:enumerator} Metal = 11 + +Apple Metal GPU. +``` + +```{cpp:enumerator} XPU = 12 + +Intel XPU GPU. +``` + +```{cpp:enumerator} MPS = 13 + +Apple Metal Performance Shaders. +``` + +```{cpp:enumerator} Meta = 14 + +Meta tensors (shape only, no data). +``` + +```{cpp:enumerator} HPU = 15 + +Habana HPU. +``` + +```{cpp:enumerator} Lazy = 17 + +Lazy tensors. +``` + +```{cpp:enumerator} IPU = 18 + +Graphcore IPU. +``` + +```{cpp:enumerator} MTIA = 19 + +Meta training and inference accelerator. +``` + +```{cpp:enumerator} PrivateUse1 = 20 + +Custom backend registered via `c10::register_privateuse1_backend()`. +``` + +Convenience constants: + +- `c10::kCPU`, `c10::kCUDA`, `c10::kHIP` +- `c10::kXLA`, `c10::kVulkan`, `c10::kMetal` +- `c10::kXPU`, `c10::kMPS`, `c10::kMeta` +- `c10::kHPU`, `c10::kLazy`, `c10::kIPU`, `c10::kMTIA` +- `c10::kPrivateUse1` diff --git a/docs/cpp/source/api/c10/guards.md b/docs/cpp/source/api/c10/guards.md new file mode 100644 index 0000000000000..fba83afbcee57 --- /dev/null +++ b/docs/cpp/source/api/c10/guards.md @@ -0,0 +1,52 @@ +--- +myst: + html_meta: + description: Device and stream guards in PyTorch C++ — RAII guards for managing current device and stream context. + keywords: PyTorch, C++, DeviceGuard, StreamGuard, RAII, device management +--- + +# Device Guards + +C10 provides device-agnostic RAII guards for managing the current device +context. These guards work across all backends (CUDA, XPU, etc.) and +automatically restore the previous device when they go out of scope. + +For backend-specific guards, see {doc}`../cuda/guards` and {doc}`../xpu/index`. + +## DeviceGuard + +```{doxygenclass} c10::DeviceGuard +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +{ + c10::DeviceGuard guard(c10::Device(c10::kCUDA, 1)); + // All operations here run on CUDA device 1 +} +// Previous device is restored +``` + +## OptionalDeviceGuard + +```{doxygenclass} c10::OptionalDeviceGuard +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +c10::OptionalDeviceGuard guard; +if (use_gpu) { + guard.reset_device(c10::Device(c10::kCUDA, 0)); +} +// Guard only restores device if it was set +``` diff --git a/docs/cpp/source/api/c10/index.md b/docs/cpp/source/api/c10/index.md new file mode 100644 index 0000000000000..dd34df87409fa --- /dev/null +++ b/docs/cpp/source/api/c10/index.md @@ -0,0 +1,66 @@ +--- +myst: + html_meta: + description: C10 core library API — fundamental types, device management, streams, and utilities used throughout PyTorch. + keywords: PyTorch, C++, c10, core, Device, Stream, types +--- + +# C10: Core Utilities + +C10 (Caffe2 + ATen = C10) is the core library that provides fundamental +utilities and data types used throughout PyTorch. It contains device +abstractions, memory management utilities, and common data structures. + +**When to use C10:** + +- When working with device-agnostic code (CPU, CUDA, XPU, etc.) +- When you need efficient array views without copying data +- When handling optional values or type-erased containers +- When writing code that needs to work across different PyTorch backends + +**Basic usage:** + +```cpp +#include +#include + +// Device abstraction +c10::Device device(c10::kCUDA, 0); +if (device.is_cuda()) { + std::cout << "Using CUDA device " << device.index() << std::endl; +} + +// Efficient array views (no copy) +std::vector sizes = {3, 4, 5}; +c10::ArrayRef sizes_ref(sizes); + +// Optional values +c10::optional maybe_dim = 2; +int64_t dim = maybe_dim.value_or(-1); +``` + +## Header Files + +- `c10/core/Device.h` - Device abstraction +- `c10/core/DeviceType.h` - Device type enumeration +- `c10/util/ArrayRef.h` - Non-owning array reference +- `c10/util/OptionalArrayRef.h` - Optional array reference +- `c10/util/Optional.h` - Optional value wrapper +- `c10/util/Half.h` - Half-precision float +- `c10/util/Exception.h` - Error checking macros (`TORCH_CHECK`, `TORCH_INTERNAL_ASSERT`, etc.) +- `c10/cuda/CUDAGuard.h` - CUDA device guards (see {doc}`../cuda/index`) +- `c10/cuda/CUDAStream.h` - CUDA stream management (see {doc}`../cuda/index`) +- `c10/xpu/XPUStream.h` - Intel XPU stream management +- `ATen/core/ivalue.h` - IValue for TorchScript interop + +## C10 Categories + +```{toctree} +:maxdepth: 1 + +device +guards +streams +types +utilities +``` diff --git a/docs/cpp/source/api/c10/streams.md b/docs/cpp/source/api/c10/streams.md new file mode 100644 index 0000000000000..e87e3d3fdad7a --- /dev/null +++ b/docs/cpp/source/api/c10/streams.md @@ -0,0 +1,34 @@ +--- +myst: + html_meta: + description: Stream API in PyTorch C++ — c10::Stream for asynchronous execution on devices. + keywords: PyTorch, C++, Stream, c10::Stream, asynchronous, device +--- + +# Streams + +`c10::Stream` is the device-agnostic base stream class. It provides a +common interface for working with streams across different backends +(CUDA, XPU, etc.). + +For backend-specific stream APIs, see {doc}`../cuda/streams` and {doc}`../xpu/index`. + +## Stream + +```{doxygenclass} c10::Stream +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +// Streams are typically obtained from backend-specific APIs +auto cuda_stream = c10::cuda::getCurrentCUDAStream(); + +// c10::Stream provides the common interface +c10::Device device = cuda_stream.device(); +c10::DeviceType type = cuda_stream.device_type(); +``` diff --git a/docs/cpp/source/api/c10/types.md b/docs/cpp/source/api/c10/types.md new file mode 100644 index 0000000000000..60ffa1d784c45 --- /dev/null +++ b/docs/cpp/source/api/c10/types.md @@ -0,0 +1,225 @@ +--- +myst: + html_meta: + description: Core types in PyTorch C++ — ArrayRef, optional, Dict, List, IListRef, Half, and IValue. + keywords: PyTorch, C++, c10, ArrayRef, optional, Dict, List, IListRef, Half, IValue +--- + +# Core Types + +C10 provides fundamental types used throughout PyTorch. + +## ArrayRef + +```{doxygenclass} c10::ArrayRef +:members: +:undoc-members: +``` + +**Example:** + +```cpp +std::vector sizes = {3, 4, 5}; +c10::ArrayRef sizes_ref(sizes); + +// Can also use initializer list +auto tensor = at::zeros({3, 4, 5}); // implicitly converts +``` + +## OptionalArrayRef + +```{doxygenclass} c10::OptionalArrayRef +:members: +:no-link: +``` + +**Example:** + +```cpp +void my_function(c10::OptionalArrayRef sizes = c10::nullopt) { + if (sizes.has_value()) { + for (auto s : sizes.value()) { + // process sizes + } + } +} +``` + +## Optional + +```{cpp:class} c10::optional + +A wrapper type that may or may not contain a value. +Similar to `std::optional`. +``` + +```{cpp:function} bool has_value() const + +Returns true if a value is present. +``` + +```{cpp:function} T& value() + +Returns the contained value. Throws if empty. +``` + +```{cpp:function} T value_or(T default_value) const + +Returns the value if present, otherwise returns the default. +``` + +**Example:** + +```cpp +c10::optional maybe_dim = c10::nullopt; + +if (maybe_dim.has_value()) { + std::cout << "Dim: " << maybe_dim.value() << std::endl; +} + +int64_t dim = maybe_dim.value_or(-1); // Returns -1 if empty +``` + +## Half + +```{cpp:class} c10::Half + +16-bit floating point type (IEEE 754 half-precision). +``` + +```{cpp:function} Half(float value) + +Construct from a float. +``` + +```{cpp:function} operator float() const + +Convert to float. +``` + +**Example:** + +```cpp +c10::Half h = 3.14f; +float f = static_cast(h); +``` + +## Containers + +C10 provides container types that store `IValue` elements internally. These +are pointer types: copies share the same underlying storage. + +### Dict + +An ordered hash map from `Key` to `Value`. Valid key types are `int64_t`, +`double`, `bool`, `std::string`, and `at::Tensor`. + +```{doxygenclass} c10::Dict +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +c10::Dict named_tensors; +named_tensors.insert("weight", torch::randn({3, 3})); +named_tensors.insert("bias", torch::zeros({3})); + +if (named_tensors.contains("weight")) { + at::Tensor w = named_tensors.at("weight"); +} + +for (const auto& entry : named_tensors) { + std::cout << entry.key() << ": " << entry.value().sizes() << std::endl; +} +``` + +### List + +A type-safe list container backed by `IValue` elements. + +```{doxygenclass} c10::List +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +c10::List tensor_list; +tensor_list.push_back(torch::randn({2, 3})); +tensor_list.push_back(torch::zeros({2, 3})); + +at::Tensor first = tensor_list.get(0); +std::cout << "List size: " << tensor_list.size() << std::endl; + +c10::List int_list; +int_list.push_back(1); +int_list.push_back(2); +int_list.push_back(3); +``` + +### IListRef + +`c10::IListRef` is a lightweight reference type that provides a unified +interface over different list-like types (`List`, `ArrayRef`, +`std::vector`). It avoids copying when passing list arguments to operators. + +```{doxygenclass} c10::IListRef +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +// IListRef can wrap different underlying types +std::vector vec = {torch::randn({2}), torch::randn({3})}; +c10::IListRef ref(vec); + +for (const auto& t : ref) { + std::cout << t.sizes() << std::endl; +} +``` + +## IValue + +`c10::IValue` (Interpreter Value) is a type-erased container used extensively +for storing values of different types. It can hold tensors, +scalars, lists, dictionaries, and other types. + +```{note} + +The full API documentation for IValue is complex due to its many type +conversion methods. See the header file `ATen/core/ivalue.h` for complete +details. +``` + +**Common methods:** + +- `isTensor()` / `toTensor()` - Check if tensor / convert to tensor +- `isInt()` / `toInt()` - Check if int / convert to int +- `isDouble()` / `toDouble()` - Check if double / convert to double +- `isBool()` / `toBool()` - Check if bool / convert to bool +- `isString()` / `toString()` - Check if string / convert to string +- `isList()` / `toList()` - Check if list / convert to list +- `isGenericDict()` / `toGenericDict()` - Check if dict / convert to dict +- `isTuple()` / `toTuple()` - Check if tuple / convert to tuple +- `isNone()` - Check if None/null + +**Example:** + +```cpp +c10::IValue val = at::ones({2, 2}); + +if (val.isTensor()) { + at::Tensor t = val.toTensor(); +} +``` diff --git a/docs/cpp/source/api/c10/utilities.md b/docs/cpp/source/api/c10/utilities.md new file mode 100644 index 0000000000000..2e61c14c8e8dc --- /dev/null +++ b/docs/cpp/source/api/c10/utilities.md @@ -0,0 +1,153 @@ +--- +myst: + html_meta: + description: C10 utility classes in PyTorch C++ — Flags, QEngine, and Reduction enumerations. + keywords: PyTorch, C++, c10, utilities, QEngine, Reduction +--- + +# Utilities + +C10 provides utility classes for memory management and other common patterns. + +## MaybeOwned + +`MaybeOwned` is a C++ smart pointer class that dynamically +encodes whether a Tensor is *owned* or *borrowed*. It is used in +certain performance-sensitive situations to avoid unnecessarily +incrementing a Tensor's reference count (at a small cost in +overhead from the extra indirection). + +```{warning} + + MaybeOwned must be used with **extreme** care. Claims of (non-)ownership + are not statically checked, and mistakes can cause reference undercounting + and use-after-free crashes. + + Due to this lack of safety net, we discourage the use of MaybeOwned + outside code paths that are known to be highly performance sensitive. + However, if you encounter pre-existing uses of MaybeOwned in code that + you want to modify, it's critical to understand how to use it correctly. +``` + +**Use Case:** + +The primary use case for `MaybeOwned` is a function or method that +dynamically chooses between returning one of its arguments (typically +from a passthrough or "no-op" code path) and returning a freshly constructed +Tensor. Such a function would return a `MaybeOwned` in both cases: +the former in a "borrowed" state via `MaybeOwned::borrowed()`, +and the latter in an "owned" state via `MaybeOwned::owned()`. + +**Example - expect_contiguous:** + +The canonical example is `Tensor`'s `expect_contiguous` method, which shortcuts +and returns a borrowed self-reference when already contiguous: + +```cpp +inline c10::MaybeOwned Tensor::expect_contiguous( + MemoryFormat memory_format) const & { + if (is_contiguous(memory_format)) { + return c10::MaybeOwned::borrowed(*this); + } else { + return c10::MaybeOwned::owned( + __dispatch_contiguous(memory_format)); + } +} +``` + +Using the vocabulary of lifetimes, the essential safety requirement for borrowing +is that a borrowed Tensor must outlive any borrowing references to it. In the example +above, we can safely borrow `*this`, but the Tensor returned by +`__dispatch_contiguous()` is freshly created, and borrowing a reference would +effectively leave it ownerless. + +**Rules of Thumb:** + +- When in doubt, don't use `MaybeOwned` at all - in particular, prefer + avoiding using it in code that doesn't use it already. New usage should only be + introduced when critical (and demonstrable) performance gains result. + +- When modifying or calling code that already uses `MaybeOwned`, remember + that it's always safe to produce a `MaybeOwned` from a Tensor in hand + via a call to `MaybeOwned::owned()`. This may result in an unnecessary + reference count, but never in misbehavior - so it's always the safer bet, unless + the lifetime of the Tensor you're looking to wrap is crystal clear. + +More details and implementation code can be found at +[MaybeOwned.h](https://github.com/pytorch/pytorch/blob/main/c10/util/MaybeOwned.h) and +[TensorBody.h](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/templates/TensorBody.h). + +## Error Handling and Assertions + +PyTorch provides macros for error checking and assertions that produce +informative error messages with source location. These are defined in +`c10/util/Exception.h`. + +### TORCH_CHECK + +The primary macro for validating user input and runtime conditions. On failure, +raises `c10::Error` (which becomes `RuntimeError` in Python). + +```cpp +#include + +// Basic check +TORCH_CHECK(tensor.dim() == 2, "Expected 2D tensor, got ", tensor.dim(), "D"); + +// Without message (default message generated) +TORCH_CHECK(x >= 0); +``` + +Typed variants raise specific Python exception types: + +- `TORCH_CHECK_INDEX(cond, ...)` — raises `IndexError` +- `TORCH_CHECK_VALUE(cond, ...)` — raises `ValueError` +- `TORCH_CHECK_TYPE(cond, ...)` — raises `TypeError` +- `TORCH_CHECK_LINALG(cond, ...)` — raises `LinAlgError` +- `TORCH_CHECK_NOT_IMPLEMENTED(cond, ...)` — raises `NotImplementedError` + +### TORCH_INTERNAL_ASSERT + +For internal invariants that should always hold (i.e., failures indicate a bug +in PyTorch, not user error). Produces a message asking users to report the bug. + +```cpp +TORCH_INTERNAL_ASSERT(googol > 0); +TORCH_INTERNAL_ASSERT(googol > 0, "googol was ", googol); +``` + +```{note} + +Use `TORCH_CHECK` for conditions that can fail due to user input. +Use `TORCH_INTERNAL_ASSERT` only for conditions that indicate a PyTorch bug. +`TORCH_INTERNAL_ASSERT_DEBUG_ONLY` is the debug-build-only variant for +hot paths. +``` + +### TORCH_WARN + +Issues a warning (not an error) to the user. + +```cpp +TORCH_WARN("This operation is slow for sparse tensors"); +TORCH_WARN_ONCE("This warning appears only once"); +``` + +### c10::Error + +The base exception class for PyTorch C++ errors. Provides source location +and optional backtrace. + +```cpp +try { + auto result = some_operation(); +} catch (const c10::Error& e) { + std::cerr << e.what() << std::endl; + // Or without backtrace: + std::cerr << e.what_without_backtrace() << std::endl; +} +``` + +Specialized subclasses: `c10::IndexError`, `c10::ValueError`, +`c10::TypeError`, `c10::NotImplementedError`, `c10::LinAlgError`, +`c10::OutOfMemoryError`. diff --git a/docs/cpp/source/api/cuda/guards.md b/docs/cpp/source/api/cuda/guards.md new file mode 100644 index 0000000000000..d7724e390525a --- /dev/null +++ b/docs/cpp/source/api/cuda/guards.md @@ -0,0 +1,96 @@ +--- +myst: + html_meta: + description: CUDA device and stream guards in PyTorch C++ — CUDAGuard, CUDAStreamGuard, and CUDAMultiStreamGuard. + keywords: PyTorch, C++, CUDA, CUDAGuard, CUDAStreamGuard, device guard, multi-GPU +--- + +# CUDA Guards + +CUDA guards are RAII wrappers that set a CUDA device or stream as the current +context and automatically restore the previous context when the guard goes +out of scope. + +## CUDAGuard + +```{doxygenstruct} c10::cuda::CUDAGuard +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +{ + c10::cuda::CUDAGuard guard(1); // Switch to device 1 + // All CUDA operations here run on device 1 + auto tensor = torch::zeros({2, 2}, torch::device(torch::kCUDA)); +} +// Previous device is restored +``` + +## CUDAStreamGuard + +```{doxygenstruct} c10::cuda::CUDAStreamGuard +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +auto stream = c10::cuda::getStreamFromPool(); +{ + c10::cuda::CUDAStreamGuard guard(stream); + // Operations here use the specified stream +} +// Previous stream is restored +``` + +## OptionalCUDAGuard + +```{doxygenstruct} c10::cuda::OptionalCUDAGuard +:members: +:undoc-members: +``` + +**Example:** + +```cpp +c10::cuda::OptionalCUDAGuard guard; +if (use_cuda) { + guard.set_device(0); +} +// Guard only switches device if set_device was called +``` + +## OptionalCUDAStreamGuard + +```{doxygenstruct} c10::cuda::OptionalCUDAStreamGuard +:members: +:undoc-members: +``` + +## CUDAMultiStreamGuard + +```{doxygenstruct} c10::cuda::CUDAMultiStreamGuard +:members: +:undoc-members: +``` + +**Example:** + +```cpp +at::cuda::CUDAStream stream0 = at::cuda::getStreamFromPool(false, 0); +at::cuda::CUDAStream stream1 = at::cuda::getStreamFromPool(false, 1); + +{ + at::cuda::CUDAMultiStreamGuard multi_guard({stream0, stream1}); + // stream0 is current on device 0, stream1 on device 1 +} +// Both streams restored +``` diff --git a/docs/cpp/source/api/cuda/index.md b/docs/cpp/source/api/cuda/index.md new file mode 100644 index 0000000000000..0f7f07f447b4f --- /dev/null +++ b/docs/cpp/source/api/cuda/index.md @@ -0,0 +1,58 @@ +--- +myst: + html_meta: + description: PyTorch CUDA C++ API — device management, streams, guards, and cuDNN/cuBLAS utilities. + keywords: PyTorch, C++, CUDA, GPU, streams, guards, cuDNN, cuBLAS +--- + +# CUDA Support + +PyTorch provides comprehensive CUDA support for GPU-accelerated tensor +operations and neural network training. The CUDA API allows you to manage +GPU devices, streams for asynchronous execution, and memory efficiently. + +**When to use CUDA APIs:** + +- When you need explicit control over which GPU device to use +- When implementing custom CUDA kernels or operations +- When optimizing performance with asynchronous stream execution +- When managing multi-GPU workloads + +**Basic usage:** + +```cpp +#include +#include + +// Check if CUDA is available +if (torch::cuda::is_available()) { + // Create tensor on GPU + auto tensor = torch::randn({2, 3}, torch::device(torch::kCUDA)); + + // Switch to a specific GPU + c10::cuda::CUDAGuard guard(0); // Use GPU 0 + + // Get the current CUDA stream + auto stream = c10::cuda::getCurrentCUDAStream(); + + // Move model to GPU + model->to(torch::kCUDA); +} +``` + +## Header Files + +- `c10/cuda/CUDAStream.h` - CUDA stream management +- `c10/cuda/CUDAGuard.h` - CUDA device guards +- `ATen/cuda/CUDAContext.h` - CUDA context management +- `ATen/cudnn/Descriptors.h` - cuDNN tensor descriptors + +## CUDA Categories + +```{toctree} +:maxdepth: 1 + +streams +guards +utilities +``` diff --git a/docs/cpp/source/api/cuda/streams.md b/docs/cpp/source/api/cuda/streams.md new file mode 100644 index 0000000000000..a55eb609d1bf7 --- /dev/null +++ b/docs/cpp/source/api/cuda/streams.md @@ -0,0 +1,190 @@ +--- +myst: + html_meta: + description: CUDA streams in PyTorch C++ — CUDAStream for asynchronous GPU execution and synchronization. + keywords: PyTorch, C++, CUDA, CUDAStream, stream, asynchronous, GPU, synchronization +--- + +# CUDA Streams + +CUDA streams provide a mechanism for asynchronous execution of operations +on the GPU. Operations queued to the same stream execute in order, while +operations on different streams can execute concurrently. + +## CUDAStream + +```{doxygenclass} c10::cuda::CUDAStream +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +// Get the default stream for current device +auto stream = c10::cuda::getDefaultCUDAStream(); + +// Create a new stream +auto new_stream = c10::cuda::getStreamFromPool(); + +// Get current stream +auto current = c10::cuda::getCurrentCUDAStream(); + +// Synchronize +stream.synchronize(); +``` + +## Acquiring CUDA Streams + +PyTorch provides several ways to acquire CUDA streams: + +1. **From the stream pool** (round-robin allocation): + + ```cpp + // Normal priority stream + at::cuda::CUDAStream stream = at::cuda::getStreamFromPool(); + + // High priority stream + at::cuda::CUDAStream high_prio = at::cuda::getStreamFromPool(/*isHighPriority=*/true); + + // Stream for specific device + at::cuda::CUDAStream dev1_stream = at::cuda::getStreamFromPool(false, /*device=*/1); + ``` + +2. **Default stream** (where most computation occurs): + + ```cpp + at::cuda::CUDAStream defaultStream = at::cuda::getDefaultCUDAStream(); + ``` + +3. **Current stream** (may differ if changed with guards): + + ```cpp + at::cuda::CUDAStream currentStream = at::cuda::getCurrentCUDAStream(); + ``` + +## Setting CUDA Streams + +**Using setCurrentCUDAStream:** + +```cpp +torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(torch::kCUDA)); + +// Get a new stream and set it as current +at::cuda::CUDAStream myStream = at::cuda::getStreamFromPool(); +at::cuda::setCurrentCUDAStream(myStream); + +// Operations now use myStream +tensor0.sum(); + +// Restore default stream +at::cuda::setCurrentCUDAStream(at::cuda::getDefaultCUDAStream()); +``` + +**Using CUDAStreamGuard (recommended):** + +```cpp +torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(torch::kCUDA)); +at::cuda::CUDAStream myStream = at::cuda::getStreamFromPool(); + +{ + at::cuda::CUDAStreamGuard guard(myStream); + // Operations use myStream within this scope + tensor0.sum(); +} +// Stream automatically restored to default +``` + +## Multi-Device Stream Management + +**Streams on multiple devices:** + +```cpp +// Acquire streams for different devices +at::cuda::CUDAStream stream0 = at::cuda::getStreamFromPool(false, 0); +at::cuda::CUDAStream stream1 = at::cuda::getStreamFromPool(false, 1); + +// Set current streams on each device +at::cuda::setCurrentCUDAStream(stream0); +at::cuda::setCurrentCUDAStream(stream1); + +// Create tensors on device 0 +torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(at::kCUDA)); +tensor0.sum(); // Uses stream0 + +// Switch to device 1 +{ + at::cuda::CUDAGuard device_guard(1); + torch::Tensor tensor1 = torch::ones({2, 2}, torch::device(at::kCUDA)); + tensor1.sum(); // Uses stream1 +} +``` + +**Using CUDAMultiStreamGuard:** + +```cpp +torch::Tensor tensor0 = torch::ones({2, 2}, torch::device({torch::kCUDA, 0})); +torch::Tensor tensor1 = torch::ones({2, 2}, torch::device({torch::kCUDA, 1})); + +at::cuda::CUDAStream stream0 = at::cuda::getStreamFromPool(false, 0); +at::cuda::CUDAStream stream1 = at::cuda::getStreamFromPool(false, 1); + +{ + // Set streams on both devices simultaneously + at::cuda::CUDAMultiStreamGuard multi_guard({stream0, stream1}); + + tensor0.sum(); // Uses stream0 on device 0 + tensor1.sum(); // Uses stream1 on device 1 +} +// Both streams restored to defaults +``` + +```{attention} + +`CUDAMultiStreamGuard` does not change the current device index. It only +changes the stream on each passed-in stream's device. +``` + +## Multi-Device Stream Handling Pattern + +The following skeleton shows three common patterns for acquiring and setting +streams across multiple CUDA devices: + +```cpp +// Create stream vectors on device 0 +std::vector streams0 = + {at::cuda::getDefaultCUDAStream(), at::cuda::getStreamFromPool()}; +at::cuda::setCurrentCUDAStream(streams0[0]); + +// Create stream vector on device 1 using CUDAGuard +std::vector streams1; +{ + at::cuda::CUDAGuard device_guard(1); + streams1.push_back(at::cuda::getDefaultCUDAStream()); + streams1.push_back(at::cuda::getStreamFromPool()); +} +at::cuda::setCurrentCUDAStream(streams1[0]); + +// Pattern 1: CUDAGuard changes current device only, not streams +{ + at::cuda::CUDAGuard device_guard(1); + // current device is 1, current stream on device 1 is still streams1[0] +} + +// Pattern 2: CUDAStreamGuard changes both current device and current stream +{ + at::cuda::CUDAStreamGuard stream_guard(streams1[1]); + // current device is 1, current stream is streams1[1] +} +// restored to device 0, stream streams0[0] + +// Pattern 3: CUDAMultiStreamGuard sets streams on multiple devices at once +{ + at::cuda::CUDAMultiStreamGuard multi_guard({streams0[1], streams1[1]}); + // current device unchanged (still 0) + // stream on device 0 is streams0[1], stream on device 1 is streams1[1] +} +// streams restored to streams0[0] and streams1[0] +``` diff --git a/docs/cpp/source/api/cuda/utilities.md b/docs/cpp/source/api/cuda/utilities.md new file mode 100644 index 0000000000000..951c49e3b71de --- /dev/null +++ b/docs/cpp/source/api/cuda/utilities.md @@ -0,0 +1,226 @@ +--- +myst: + html_meta: + description: CUDA utility functions in PyTorch C++ — device properties, cuBLAS/cuDNN handles, and stream management. + keywords: PyTorch, C++, CUDA, device_count, cuBLAS, cuDNN, cuSPARSE, TensorDescriptor +--- + +# CUDA Utility Functions + +PyTorch provides utility functions for querying and managing CUDA devices, +streams, and library handles. + +## Device Management + +```{doxygenfunction} torch::cuda::device_count +``` + +```{cpp:function} int c10::cuda::current_device() + +Returns the index of the current CUDA device. +``` + +**Example:** + +```cpp +#include + +// Check available devices +int num_devices = c10::cuda::device_count(); + +// Get current device +int current = c10::cuda::current_device(); +``` + +## Device Properties + +```{doxygenfunction} at::cuda::getCurrentDeviceProperties +``` + +```{doxygenfunction} at::cuda::getDeviceProperties +``` + +```{doxygenfunction} at::cuda::canDeviceAccessPeer +``` + +```{doxygenfunction} at::cuda::warp_size +``` + +**Example:** + +```cpp +#include + +// Query properties of the current device +cudaDeviceProp* props = at::cuda::getCurrentDeviceProperties(); +std::cout << "Device: " << props->name << std::endl; +std::cout << "Compute capability: " << props->major << "." << props->minor << std::endl; + +// Query a specific device +cudaDeviceProp* dev1_props = at::cuda::getDeviceProperties(1); + +// Check peer access +bool can_access = at::cuda::canDeviceAccessPeer(0, 1); +``` + +## Library Handles + +These functions return handles for CUDA math libraries on the current device +and stream. They are primarily useful when writing custom CUDA kernels that +call cuBLAS, cuSPARSE, or cuSOLVER directly. + +```{doxygenfunction} at::cuda::getCurrentCUDABlasHandle +``` + +```{doxygenfunction} at::cuda::getCurrentCUDABlasLtHandle +``` + +```{doxygenfunction} at::cuda::getCurrentCUDASparseHandle +``` + +```{doxygenfunction} at::cuda::getCurrentCUDASolverDnHandle +``` + +**Example:** + +```cpp +#include + +// Get cuBLAS handle for current device/stream +cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle(); + +// Get cuSPARSE handle +cusparseHandle_t sparse_handle = at::cuda::getCurrentCUDASparseHandle(); +``` + +## cuDNN Descriptors + +When writing custom kernels that call cuDNN directly, PyTorch provides RAII +wrapper classes for cuDNN descriptors. These are defined in +`ATen/cudnn/Descriptors.h`. + +### Descriptor (Base Class) + +```{doxygenclass} at::native::Descriptor +:members: +:undoc-members: +``` + +A generic RAII wrapper for cuDNN descriptor types. Descriptors default +construct to a `nullptr` and are initialized on first use via `mut_desc()`. +Use `desc()` for read-only access. + +### TensorDescriptor + +```{doxygenclass} at::native::TensorDescriptor +:members: +:undoc-members: +``` + +Wraps `cudnnTensorDescriptor_t`. Supports padding lower-dimensional tensors +to meet cuDNN broadcasting requirements (see `pad` parameter). + +**Example:** + +```cpp +#include + +at::Tensor input = torch::randn({32, 3, 224, 224}, torch::kCUDA); +at::native::TensorDescriptor desc(input); +cudnnTensorDescriptor_t raw = desc.desc(); +``` + +### FilterDescriptor + +```{doxygenclass} at::native::FilterDescriptor +:members: +:undoc-members: +``` + +Wraps `cudnnFilterDescriptor_t` for convolution filter weights. + +### ConvolutionDescriptor + +```{doxygenstruct} at::native::ConvolutionDescriptor +:members: +:undoc-members: +``` + +Wraps `cudnnConvolutionDescriptor_t`. Configures padding, stride, dilation, +groups, and math type (TF32, tensor ops) for convolution operations. + +### RNNDataDescriptor + +```{doxygenclass} at::native::RNNDataDescriptor +:members: +:undoc-members: +``` + +Wraps `cudnnRNNDataDescriptor_t` for variable-length sequence data. + +### DropoutDescriptor + +```{doxygenstruct} at::native::DropoutDescriptor +:members: +:undoc-members: +``` + +Wraps `cudnnDropoutDescriptor_t`. Manages RNG state for cuDNN dropout. + +### ActivationDescriptor + +```{doxygenstruct} at::native::ActivationDescriptor +:members: +:undoc-members: +``` + +Wraps `cudnnActivationDescriptor_t`. + +### SpatialTransformerDescriptor + +```{doxygenstruct} at::native::SpatialTransformerDescriptor +:members: +:undoc-members: +``` + +### CTCLossDescriptor + +```{doxygenstruct} at::native::CTCLossDescriptor +:members: +:undoc-members: +``` + +## Stream Management + +```{doxygenfunction} c10::cuda::getDefaultCUDAStream +``` + +```{doxygenfunction} c10::cuda::getCurrentCUDAStream +``` + +```{doxygenfunction} c10::cuda::setCurrentCUDAStream +``` + +```{doxygenfunction} c10::cuda::getStreamFromPool(const bool isHighPriority, DeviceIndex device) +``` + +```{doxygenfunction} c10::cuda::getStreamFromExternal +``` + +**Example:** + +```cpp +#include + +// Create and set custom stream +auto stream = c10::cuda::getStreamFromPool(); +c10::cuda::setCurrentCUDAStream(stream); + +// Get default stream +auto default_stream = c10::cuda::getDefaultCUDAStream(); + +// Wrap an externally created CUDA stream +cudaStream_t ext_stream; +cudaStreamCreate(&ext_stream); +auto wrapped = c10::cuda::getStreamFromExternal(ext_stream, /*device_index=*/0); +``` diff --git a/docs/cpp/source/api/data/dataloader.md b/docs/cpp/source/api/data/dataloader.md new file mode 100644 index 0000000000000..f368363368bf6 --- /dev/null +++ b/docs/cpp/source/api/data/dataloader.md @@ -0,0 +1,103 @@ +--- +myst: + html_meta: + description: DataLoader in PyTorch C++ — parallel data loading with batching, sampling, and multi-worker support. + keywords: PyTorch, C++, DataLoader, data loading, batching, workers, make_data_loader +--- + +# DataLoader + +The DataLoader batches samples from a dataset and optionally shuffles and +parallelizes the loading process. It is the main interface for iterating +over training data. + +## DataLoader Classes + +```{doxygenclass} torch::data::DataLoaderBase +:members: +:undoc-members: +``` + +## DataLoaderOptions + +```{doxygenstruct} torch::data::DataLoaderOptions +:members: +:undoc-members: +``` + +## StatefulDataLoader + +A DataLoader for `StatefulDataset` types that manage their own batching logic +internally. + +```{doxygenclass} torch::data::StatefulDataLoader +:members: +:undoc-members: +``` + +## StatelessDataLoader + +A DataLoader for `Dataset` types that use external samplers for batching. + +```{doxygenclass} torch::data::StatelessDataLoader +:members: +:undoc-members: +``` + +## Iterator + +```{doxygenclass} torch::data::Iterator +:members: +:undoc-members: +``` + +## Creating a DataLoader + +Use `make_data_loader` to create a DataLoader from a dataset: + +```cpp +auto data_loader = torch::data::make_data_loader( + std::move(dataset), + torch::data::DataLoaderOptions() + .batch_size(64) + .workers(4)); + +for (auto& batch : *data_loader) { + auto data = batch.data; + auto target = batch.target; + // Train on batch +} +``` + +## Complete Training Example + +```cpp +#include + +int main() { + // Load dataset + auto dataset = torch::data::datasets::MNIST("./data") + .map(torch::data::transforms::Normalize<>(0.1307, 0.3081)) + .map(torch::data::transforms::Stack<>()); + + // Create data loader + auto data_loader = torch::data::make_data_loader( + std::move(dataset), + torch::data::DataLoaderOptions().batch_size(64).workers(2)); + + // Create model and optimizer + auto model = std::make_shared(); + auto optimizer = torch::optim::Adam(model->parameters(), 0.001); + + // Training loop + for (size_t epoch = 1; epoch <= 10; ++epoch) { + for (auto& batch : *data_loader) { + optimizer.zero_grad(); + auto output = model->forward(batch.data); + auto loss = torch::nll_loss(output, batch.target); + loss.backward(); + optimizer.step(); + } + } +} +``` diff --git a/docs/cpp/source/api/data/datasets.md b/docs/cpp/source/api/data/datasets.md new file mode 100644 index 0000000000000..4bb01e255c393 --- /dev/null +++ b/docs/cpp/source/api/data/datasets.md @@ -0,0 +1,110 @@ +--- +myst: + html_meta: + description: Dataset classes in PyTorch C++ — Dataset, MapDataset, StreamDataset, and built-in datasets like MNIST. + keywords: PyTorch, C++, Dataset, MapDataset, StreamDataset, MNIST, data +--- + +# Datasets + +The dataset abstraction defines how to access individual samples in your data. +All datasets inherit from `Dataset` and must implement `get()` and `size()`. + +## Dataset Base Class + +```{doxygenclass} torch::data::datasets::Dataset +:members: +:undoc-members: +``` + +```{doxygenclass} torch::data::datasets::BatchDataset +:members: +:undoc-members: +``` + +## StatefulDataset + +A dataset that manages its own state across batches (e.g., position in a stream). +Unlike `Dataset`, it produces batches directly without external samplers. + +```{doxygenclass} torch::data::datasets::StatefulDataset +:members: +:undoc-members: +``` + +## ChunkDataReader + +Interface for reading chunks of data from a data source. Used with +`ChunkDataset` for large-scale data loading. + +```{doxygenclass} torch::data::datasets::ChunkDataReader +:members: +:undoc-members: +``` + +## Custom Dataset Example + +```cpp +class CustomDataset : public torch::data::datasets::Dataset { + public: + explicit CustomDataset(const std::string& root) { + // Load data from root directory + } + + torch::data::Example<> get(size_t index) override { + return {images_[index], labels_[index]}; + } + + torch::optional size() const override { + return images_.size(0); + } + + private: + torch::Tensor images_, labels_; +}; +``` + +## MapDataset + +```{doxygenclass} torch::data::datasets::MapDataset +:members: +:undoc-members: +``` + +## ChunkDataset + +```{doxygenclass} torch::data::datasets::ChunkDataset +:members: +:undoc-members: +``` + +## SharedBatchDataset + +```{doxygenclass} torch::data::datasets::SharedBatchDataset +:members: +:undoc-members: +``` + +## Built-in Datasets + +### MNIST + +```{doxygenclass} torch::data::datasets::MNIST +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto dataset = torch::data::datasets::MNIST("./data") + .map(torch::data::transforms::Normalize<>(0.1307, 0.3081)) + .map(torch::data::transforms::Stack<>()); +``` + +## Example Struct + +```{doxygenstruct} torch::data::Example +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/data/index.md b/docs/cpp/source/api/data/index.md new file mode 100644 index 0000000000000..7e5c5517c510c --- /dev/null +++ b/docs/cpp/source/api/data/index.md @@ -0,0 +1,66 @@ +--- +myst: + html_meta: + description: PyTorch C++ data loading API — datasets, data loaders, samplers, and transforms. + keywords: PyTorch, C++, data, DataLoader, Dataset, sampler, transform +--- + +# Data Loading (torch::data) + +The `torch::data` namespace provides utilities for loading and processing +datasets during training. It includes dataset abstractions, data loaders for +batching and shuffling, samplers for controlling data access patterns, and +transforms for data augmentation. + +**When to use torch::data:** + +- When loading training data in batches +- When you need parallel data loading with multiple workers +- When implementing custom datasets or transforms + +**Components overview:** + +- **Dataset**: Defines how to access individual samples (implement `get()` and `size()`) +- **DataLoader**: Batches samples and optionally shuffles/parallelizes loading +- **Sampler**: Controls the order in which samples are accessed +- **Transform**: Applies preprocessing (normalization, augmentation) to samples + +**Basic usage:** + +```cpp +#include + +// Load built-in dataset +auto dataset = torch::data::datasets::MNIST("./data") + .map(torch::data::transforms::Normalize<>(0.1307, 0.3081)) + .map(torch::data::transforms::Stack<>()); + +// Create data loader with batching and shuffling +auto data_loader = torch::data::make_data_loader( + std::move(dataset), + torch::data::DataLoaderOptions().batch_size(64).workers(4)); + +// Iterate over batches +for (auto& batch : *data_loader) { + auto images = batch.data; // Shape: [64, 1, 28, 28] + auto labels = batch.target; // Shape: [64] +} +``` + +## Header Files + +- `torch/csrc/api/include/torch/data.h` - Main data header +- `torch/csrc/api/include/torch/data/dataloader.h` - DataLoader +- `torch/csrc/api/include/torch/data/datasets.h` - Dataset classes +- `torch/csrc/api/include/torch/data/samplers.h` - Samplers + +## Module Categories + +```{toctree} +:maxdepth: 1 + +datasets +dataloader +samplers +transforms +``` diff --git a/docs/cpp/source/api/data/samplers.md b/docs/cpp/source/api/data/samplers.md new file mode 100644 index 0000000000000..f18f09c27d493 --- /dev/null +++ b/docs/cpp/source/api/data/samplers.md @@ -0,0 +1,69 @@ +--- +myst: + html_meta: + description: Data samplers in PyTorch C++ — RandomSampler, SequentialSampler, DistributedRandomSampler, and StreamSampler. + keywords: PyTorch, C++, sampler, RandomSampler, SequentialSampler, DistributedRandomSampler +--- + +# Samplers + +Samplers control the order in which samples are accessed from a dataset. +They determine the indices that the DataLoader uses to fetch data. + +## Sampler Base Class + +```{doxygenclass} torch::data::samplers::Sampler +:members: +:undoc-members: +``` + +## Sequential Sampler + +Accesses samples in order from 0 to N-1. Use this for evaluation or when +order matters. + +```{doxygenclass} torch::data::samplers::SequentialSampler +:members: +:undoc-members: +``` + +## Random Sampler + +Accesses samples in random order. Use this for training to ensure the model +sees samples in different orders each epoch. + +```{doxygenclass} torch::data::samplers::RandomSampler +:members: +:undoc-members: +``` + +## Distributed Random Sampler + +For distributed training, ensures each process gets a different subset of +the data without overlap. + +```{doxygenclass} torch::data::samplers::DistributedRandomSampler +:members: +:undoc-members: +``` + +## Distributed Sampler (Base) + +```{doxygenclass} torch::data::samplers::DistributedSampler +:members: +:undoc-members: +``` + +## Distributed Sequential Sampler + +```{doxygenclass} torch::data::samplers::DistributedSequentialSampler +:members: +:undoc-members: +``` + +## Stream Sampler + +```{doxygenclass} torch::data::samplers::StreamSampler +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/data/transforms.md b/docs/cpp/source/api/data/transforms.md new file mode 100644 index 0000000000000..5026596e8b63b --- /dev/null +++ b/docs/cpp/source/api/data/transforms.md @@ -0,0 +1,95 @@ +--- +myst: + html_meta: + description: Data transforms in PyTorch C++ — Stack, Normalize, Lambda, and Collate transforms for data pipelines. + keywords: PyTorch, C++, transforms, Stack, Normalize, Lambda, Collate, data pipeline +--- + +# Transforms + +Transforms apply preprocessing to data samples, such as normalization or +augmentation. They can be chained using the `.map()` method on datasets. + +## Transform (Base Class) + +The base class for all transforms. Subclass this to create custom transforms. + +```{doxygenclass} torch::data::transforms::Transform +:members: +:undoc-members: +``` + +## BatchTransform (Base Class) + +Base class for transforms that operate on entire batches. + +```{doxygenclass} torch::data::transforms::BatchTransform +:members: +:undoc-members: +``` + +## TensorTransform + +Base class for transforms that operate on tensors specifically. + +```{doxygenclass} torch::data::transforms::TensorTransform +:members: +:undoc-members: +``` + +## Normalize + +Normalizes tensors with a given mean and standard deviation. + +```{doxygenstruct} torch::data::transforms::Normalize +:members: +:undoc-members: +``` + +## Stack + +Stacks a batch of tensors into a single tensor. + +```{doxygenstruct} torch::data::transforms::Stack +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto dataset = torch::data::datasets::MNIST("./data") + .map(torch::data::transforms::Normalize<>(0.5, 0.5)) + .map(torch::data::transforms::Stack<>()); +``` + +## Lambda + +```{doxygenclass} torch::data::transforms::Lambda +:members: +:undoc-members: +``` + +## TensorLambda + +```{doxygenclass} torch::data::transforms::TensorLambda +:members: +:undoc-members: +``` + +## BatchLambda + +```{doxygenclass} torch::data::transforms::BatchLambda +:members: +:undoc-members: +``` + +## Chaining Transforms + +Transforms can be chained together using `.map()`: + +```cpp +auto dataset = torch::data::datasets::MNIST("./data") + .map(torch::data::transforms::Normalize<>(0.1307, 0.3081)) + .map(torch::data::transforms::Stack<>()); +``` diff --git a/docs/cpp/source/api/index.md b/docs/cpp/source/api/index.md new file mode 100644 index 0000000000000..3c3ff9d855039 --- /dev/null +++ b/docs/cpp/source/api/index.md @@ -0,0 +1,40 @@ +--- +myst: + html_meta: + description: PyTorch C++ API reference — complete documentation for ATen, Autograd, nn, optim, data, CUDA, and more. + keywords: PyTorch, C++, API reference, ATen, Autograd, nn, optim, CUDA +--- + +# C++ API Reference + +This section provides reference documentation for the PyTorch C++ API, +organized by module. + +```{toctree} +:maxdepth: 2 +:caption: Core + +aten/index +c10/index +autograd/index +cuda/index +xpu/index +``` + +```{toctree} +:maxdepth: 2 +:caption: C++ Frontend + +nn/index +optim/index +data/index +serialize/index +``` + +```{toctree} +:maxdepth: 2 +:caption: Extensions + +library/index +stable/index +``` diff --git a/docs/cpp/source/api/library/custom_classes.md b/docs/cpp/source/api/library/custom_classes.md new file mode 100644 index 0000000000000..48828eb731aba --- /dev/null +++ b/docs/cpp/source/api/library/custom_classes.md @@ -0,0 +1,94 @@ +--- +myst: + html_meta: + description: Custom classes in PyTorch C++ — registering C++ classes for use in TorchScript and Python. + keywords: PyTorch, C++, custom class, TorchScript, TORCH_CLASS, registration +--- + +# Custom Classes + +PyTorch allows registering custom C++ classes that can be used from Python +and TorchScript. + +Header: `torch/custom_class.h` + +## class\_ Template + +```{doxygenclass} torch::class_ +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +struct MyClass : torch::CustomClassHolder { + int value; + + MyClass(int v) : value(v) {} + + int getValue() const { return value; } + void setValue(int v) { value = v; } +}; + +TORCH_LIBRARY(my_classes, m) { + m.class_("MyClass") + .def(torch::init()) + .def("getValue", &MyClass::getValue) + .def("setValue", &MyClass::setValue) + .def_readwrite("value", &MyClass::value); +} +``` + +## Registering Methods + +**Constructor:** + +```cpp +m.class_("MyClass") + .def(torch::init()) // Constructor taking int +``` + +**Methods:** + +```cpp +m.class_("MyClass") + .def("getValue", &MyClass::getValue) + .def("setValue", &MyClass::setValue) +``` + +**Properties:** + +```cpp +m.class_("MyClass") + .def_readwrite("value", &MyClass::value) // Read-write + .def_readonly("const_value", &MyClass::const_value) // Read-only +``` + +## Using Custom Classes + +**From C++:** + +```cpp +auto my_obj = c10::make_intrusive(42); +int val = my_obj->getValue(); +``` + +**From Python:** + +```python +import torch +torch.classes.load_library("path/to/library.so") +obj = torch.classes.my_classes.MyClass(42) +print(obj.getValue()) +``` + +**In TorchScript:** + +```python +@torch.jit.script +def use_my_class(x: torch.classes.my_classes.MyClass) -> int: + return x.getValue() +``` diff --git a/docs/cpp/source/api/library/index.md b/docs/cpp/source/api/library/index.md new file mode 100644 index 0000000000000..2d3e1e225a1be --- /dev/null +++ b/docs/cpp/source/api/library/index.md @@ -0,0 +1,62 @@ +--- +myst: + html_meta: + description: PyTorch Library API in C++ — operator registration, custom classes, and versioning. + keywords: PyTorch, C++, Library, operator registration, custom class, dispatch +--- + +# Torch Library API + +The Torch Library API provides capabilities for extending PyTorch's core library +of operators with user-defined operators and data types. This is the primary +mechanism for registering custom C++ operators that can be called from both +Python and C++. + +**When to use the Library API:** + +- When creating custom operators for PyTorch +- When implementing backend-specific kernels (CPU, CUDA, etc.) +- When registering custom classes for use in TorchScript +- When extending PyTorch with new functionality + +**Basic usage:** + +```cpp +#include + +// Define a custom operator +torch::Tensor my_add(const torch::Tensor& a, const torch::Tensor& b) { + return a + b; +} + +// Register the operator +TORCH_LIBRARY(myops, m) { + m.def("add(Tensor a, Tensor b) -> Tensor", &my_add); +} + +// Use from C++ +auto result = torch::dispatcher::call("myops::add", tensor_a, tensor_b); +``` + +For a tutorial-style introduction to the library API, check out the +[Extending TorchScript with Custom C++ Operators](https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html) +tutorial. + +## Header Files + +- `torch/library.h` - Main library API header +- `torch/custom_class.h` - Custom class registration + +## Library API Categories + +```{toctree} +:maxdepth: 1 + +registration +custom_classes +versioning +``` + +## See Also + +- {doc}`../stable/index` - For stable ABI operator registration diff --git a/docs/cpp/source/api/library/registration.md b/docs/cpp/source/api/library/registration.md new file mode 100644 index 0000000000000..ec74871915681 --- /dev/null +++ b/docs/cpp/source/api/library/registration.md @@ -0,0 +1,113 @@ +--- +myst: + html_meta: + description: Operator registration in PyTorch C++ — TORCH_LIBRARY, TORCH_LIBRARY_IMPL for custom operators. + keywords: PyTorch, C++, TORCH_LIBRARY, operator registration, dispatch, custom op +--- + +# Operator Registration + +The library API provides macros and classes for registering custom operators +with PyTorch's dispatcher. + +## Macros + +### TORCH_LIBRARY + +```{doxygendefine} TORCH_LIBRARY +``` + +**Example:** + +```cpp +TORCH_LIBRARY(myops, m) { + m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl); + m.def("mul(Tensor self, Tensor other) -> Tensor"); + m.impl("mul", torch::kCPU, &mul_cpu_impl); + m.impl("mul", torch::kCUDA, &mul_cuda_impl); +} +``` + +### TORCH_LIBRARY_IMPL + +```{doxygendefine} TORCH_LIBRARY_IMPL +``` + +**Example:** + +```cpp +TORCH_LIBRARY_IMPL(myops, XLA, m) { + m.impl("mul", &mul_xla_impl); +} +``` + +### TORCH_LIBRARY_FRAGMENT + +```{doxygendefine} TORCH_LIBRARY_FRAGMENT +``` + +**Example:** + +```cpp +// In file1.cpp +TORCH_LIBRARY(myops, m) { + m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl); +} + +// In file2.cpp +TORCH_LIBRARY_FRAGMENT(myops, m) { + m.def("mul(Tensor self, Tensor other) -> Tensor", &mul_impl); +} +``` + +## Classes + +### Library + +```{doxygenclass} torch::Library +``` + +**Example:** + +```cpp +TORCH_LIBRARY(myops, m) { + // Define with implementation + m.def("add(Tensor self, Tensor other) -> Tensor", &add_impl); + + // Define schema only + m.def("mul(Tensor self, Tensor other) -> Tensor"); + + // Provide backend-specific implementations + m.impl("mul", torch::kCPU, &mul_cpu_impl); + m.impl("mul", torch::kCUDA, &mul_cuda_impl); +} +``` + +### CppFunction + +```{doxygenclass} torch::CppFunction +:members: +:no-link: +``` + +### OrderedDict + +```{doxygenclass} torch::OrderedDict +:members: +:undoc-members: +``` + +## Functions + +The library API provides builder methods on the `Library` class for registering +operators. See the `Library` class documentation above for the full API including +`def()`, `impl()`, and `fallback()` methods. + +## Dispatch Keys + +Common dispatch keys used with `torch::dispatch()`: + +- `torch::kCPU` - CPU backend +- `torch::kCUDA` - CUDA backend +- `torch::kAutograd` - Autograd backend +- `torch::kMeta` - Meta tensor backend diff --git a/docs/cpp/source/api/library/versioning.md b/docs/cpp/source/api/library/versioning.md new file mode 100644 index 0000000000000..a287f3c2a5751 --- /dev/null +++ b/docs/cpp/source/api/library/versioning.md @@ -0,0 +1,44 @@ +--- +myst: + html_meta: + description: Operator versioning in PyTorch C++ — managing backward compatibility for serialized models. + keywords: PyTorch, C++, versioning, backward compatibility, operator, serialization +--- + +# Library Versioning + +PyTorch provides version number macros for identifying the version of LibTorch in use. + +**Example:** + +```cpp +#include +#include + +int main() { + std::cout << "PyTorch version from parts: " + << TORCH_VERSION_MAJOR << "." + << TORCH_VERSION_MINOR << "." + << TORCH_VERSION_PATCH << std::endl; + std::cout << "PyTorch version: " << TORCH_VERSION << std::endl; +} +``` + +This will output something like: + +```text +PyTorch version from parts: 1.8.0 +PyTorch version: 1.8.0 +``` + +```{note} + +These macros are only available in PyTorch >= 1.8.0. +``` + +## Version Macros + +- `TORCH_VERSION_MAJOR` - Major version number +- `TORCH_VERSION_MINOR` - Minor version number +- `TORCH_VERSION_PATCH` - Patch version number +- `TORCH_VERSION` - Full version string (e.g., "1.8.0") diff --git a/docs/cpp/source/api/nn/activation.md b/docs/cpp/source/api/nn/activation.md new file mode 100644 index 0000000000000..d4551490f2f70 --- /dev/null +++ b/docs/cpp/source/api/nn/activation.md @@ -0,0 +1,356 @@ +--- +myst: + html_meta: + description: Activation functions in PyTorch C++ — ReLU, GELU, Sigmoid, Softmax, and more torch::nn activation modules. + keywords: PyTorch, C++, activation, ReLU, GELU, Sigmoid, Softmax, LeakyReLU, ELU, Mish +--- + +# Activation Functions + +Activation functions introduce non-linearity into neural networks, allowing them +to learn complex patterns. Without activations, stacked linear layers would collapse +into a single linear transformation. + +**Common choices:** + +- **ReLU family** (ReLU, LeakyReLU, PReLU, RReLU): Fast, widely used, good default choice +- **ELU family** (ELU, SELU, CELU): Smoother than ReLU, can produce negative outputs +- **GELU/SiLU/Mish**: Modern activations popular in transformers and advanced architectures +- **Sigmoid/Tanh**: Classic activations, useful for output layers (probabilities, bounded outputs) +- **Softmax**: Converts logits to probability distribution (classification output) + +## ReLU + +```{doxygenclass} torch::nn::ReLU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ReLUImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto relu = torch::nn::ReLU(torch::nn::ReLUOptions().inplace(true)); +``` + +## LeakyReLU + +```{doxygenclass} torch::nn::LeakyReLU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LeakyReLUImpl +:members: +:undoc-members: +``` + +## PReLU + +```{doxygenclass} torch::nn::PReLU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::PReLUImpl +:members: +:undoc-members: +``` + +## RReLU + +```{doxygenclass} torch::nn::RReLU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::RReLUImpl +:members: +:undoc-members: +``` + +## ReLU6 + +Like ReLU but caps the output at 6: `min(max(0, x), 6)`. Commonly used in +mobile architectures (MobileNet). + +```{doxygenclass} torch::nn::ReLU6 +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ReLU6Impl +:members: +:undoc-members: +``` + +## GLU + +Gated Linear Unit. Splits the input tensor in half along a dimension, +then applies `a * sigmoid(b)`. + +```{doxygenclass} torch::nn::GLU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::GLUImpl +:members: +:undoc-members: +``` + +## LogSigmoid + +Applies element-wise `log(sigmoid(x))`. Numerically more stable than +computing `log` and `sigmoid` separately. + +```{doxygenclass} torch::nn::LogSigmoid +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LogSigmoidImpl +:members: +:undoc-members: +``` + +## ELU + +```{doxygenclass} torch::nn::ELU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ELUImpl +:members: +:undoc-members: +``` + +## SELU + +```{doxygenclass} torch::nn::SELU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::SELUImpl +:members: +:undoc-members: +``` + +## CELU + +```{doxygenclass} torch::nn::CELU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::CELUImpl +:members: +:undoc-members: +``` + +## GELU + +```{doxygenclass} torch::nn::GELU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::GELUImpl +:members: +:undoc-members: +``` + +## SiLU (Swish) + +```{doxygenclass} torch::nn::SiLU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::SiLUImpl +:members: +:undoc-members: +``` + +## Mish + +```{doxygenclass} torch::nn::Mish +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MishImpl +:members: +:undoc-members: +``` + +## Sigmoid + +```{doxygenclass} torch::nn::Sigmoid +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::SigmoidImpl +:members: +:undoc-members: +``` + +## Tanh + +```{doxygenclass} torch::nn::Tanh +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::TanhImpl +:members: +:undoc-members: +``` + +## Softmax + +```{doxygenclass} torch::nn::Softmax +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::SoftmaxImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto softmax = torch::nn::Softmax(torch::nn::SoftmaxOptions(/*dim=*/1)); +``` + +## Softmax2d + +Applies `Softmax` over features to each spatial location in a 4D input +tensor of shape `(N, C, H, W)`. + +```{doxygenclass} torch::nn::Softmax2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::Softmax2dImpl +:members: +:undoc-members: +``` + +## LogSoftmax + +```{doxygenclass} torch::nn::LogSoftmax +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LogSoftmaxImpl +:members: +:undoc-members: +``` + +## Softmin + +```{doxygenclass} torch::nn::Softmin +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::SoftminImpl +:members: +:undoc-members: +``` + +## Softplus + +```{doxygenclass} torch::nn::Softplus +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::SoftplusImpl +:members: +:undoc-members: +``` + +## Softshrink + +```{doxygenclass} torch::nn::Softshrink +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::SoftshrinkImpl +:members: +:undoc-members: +``` + +## Softsign + +```{doxygenclass} torch::nn::Softsign +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::SoftsignImpl +:members: +:undoc-members: +``` + +## Hardshrink + +```{doxygenclass} torch::nn::Hardshrink +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::HardshrinkImpl +:members: +:undoc-members: +``` + +## Hardtanh + +```{doxygenclass} torch::nn::Hardtanh +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::HardtanhImpl +:members: +:undoc-members: +``` + +## Tanhshrink + +```{doxygenclass} torch::nn::Tanhshrink +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::TanhshrinkImpl +:members: +:undoc-members: +``` + +## Threshold + +```{doxygenclass} torch::nn::Threshold +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ThresholdImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/containers.md b/docs/cpp/source/api/nn/containers.md new file mode 100644 index 0000000000000..c9a694a96b158 --- /dev/null +++ b/docs/cpp/source/api/nn/containers.md @@ -0,0 +1,110 @@ +--- +myst: + html_meta: + description: Module containers in PyTorch C++ — Sequential, ModuleList, ModuleDict for composing neural networks. + keywords: PyTorch, C++, Sequential, ModuleList, ModuleDict, container, module +--- + +# Containers + +Container modules hold other modules and define how they are composed together. +Use containers to build complex architectures from simpler building blocks. + +- **Sequential**: Chain modules in order, output of one feeds into the next +- **ModuleList**: Store modules in a list for iteration (not auto-forwarded) +- **ModuleDict**: Store modules in a dictionary for named access +- **ParameterList/ParameterDict**: Store parameters directly without wrapping in modules + +```{note} + +PyTorch's C++ API uses the PIMPL (Pointer to Implementation) pattern. You create +modules using the public class name (e.g., `torch::nn::Sequential`), which +internally wraps an implementation class (`SequentialImpl`). The documentation +below shows the implementation classes, which contain all the actual methods. +``` + +## Sequential + +`Sequential` is a container that chains modules together. Each module's output +becomes the next module's input. This is the simplest way to build feed-forward +networks. + +```{doxygenclass} torch::nn::Sequential +:members: +:undoc-members: +``` + +**Example:** + +```cpp +torch::nn::Sequential seq( + torch::nn::Conv2d(torch::nn::Conv2dOptions(1, 32, 3)), + torch::nn::ReLU(), + torch::nn::Conv2d(torch::nn::Conv2dOptions(32, 64, 3)), + torch::nn::ReLU() +); + +auto output = seq->forward(input); +``` + +## ModuleList + +`ModuleList` stores modules in a list for indexed or iterated access. Unlike +`Sequential`, it does not have a built-in `forward()` method—you control how +modules are called. + +```{doxygenclass} torch::nn::ModuleList +:members: +:undoc-members: +``` + +**Example:** + +```cpp +torch::nn::ModuleList layers; +layers->push_back(torch::nn::Linear(10, 20)); +layers->push_back(torch::nn::Linear(20, 30)); + +torch::Tensor x = input; +for (const auto& layer : *layers) { + x = layer->as()->forward(x); +} +``` + +## ModuleDict + +`ModuleDict` stores modules in a dictionary for named access. Useful when you +need to select modules by name at runtime. + +```{doxygenclass} torch::nn::ModuleDict +:members: +:undoc-members: +``` + +## ParameterList + +`ParameterList` stores parameters directly without wrapping them in modules. + +```{doxygenclass} torch::nn::ParameterList +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ParameterListImpl +:members: +:undoc-members: +``` + +## ParameterDict + +`ParameterDict` stores parameters in a dictionary for named access. + +```{doxygenclass} torch::nn::ParameterDict +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ParameterDictImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/convolution.md b/docs/cpp/source/api/nn/convolution.md new file mode 100644 index 0000000000000..2c2addf3ce266 --- /dev/null +++ b/docs/cpp/source/api/nn/convolution.md @@ -0,0 +1,134 @@ +--- +myst: + html_meta: + description: Convolution layers in PyTorch C++ — Conv1d, Conv2d, Conv3d, and transposed convolutions. + keywords: PyTorch, C++, convolution, Conv1d, Conv2d, Conv3d, ConvTranspose, nn +--- + +# Convolution Layers + +Convolutional layers apply learnable filters to input data, extracting local features +through sliding window operations. They are fundamental to CNNs for image, audio, and +sequential data processing. + +- **Conv1d/2d/3d**: Standard convolution for 1D sequences, 2D images, or 3D volumes +- **ConvTranspose1d/2d/3d**: Transposed convolution (deconvolution) for upsampling + +**Key parameters:** + +- `in_channels`: Number of input channels +- `out_channels`: Number of output channels (number of filters) +- `kernel_size`: Size of the convolving kernel +- `stride`: Stride of the convolution (default: 1) +- `padding`: Zero-padding added to input (default: 0) +- `dilation`: Spacing between kernel elements (default: 1) +- `groups`: Number of blocked connections (default: 1, use `in_channels` for depthwise) + +## Conv1d + +Applies 1D convolution over an input signal composed of several input planes. + +```{doxygenclass} torch::nn::Conv1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::Conv1dImpl +:members: +:undoc-members: +``` + +## Conv2d + +Applies 2D convolution over an input image. The most commonly used layer for +image processing tasks. + +```{doxygenclass} torch::nn::Conv2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::Conv2dImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +// Create Conv2d: 3 input channels, 64 output channels, 3x3 kernel +auto conv = torch::nn::Conv2d( + torch::nn::Conv2dOptions(3, 64, 3) + .stride(1) + .padding(1) + .bias(true)); + +auto output = conv->forward(input); // input: [N, 3, H, W] +``` + +## Conv3d + +Applies 3D convolution over an input volume (e.g., video frames or 3D medical images). + +```{doxygenclass} torch::nn::Conv3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::Conv3dImpl +:members: +:undoc-members: +``` + +## ConvTranspose1d + +Applies 1D transposed convolution (fractionally-strided convolution) for upsampling. + +```{doxygenclass} torch::nn::ConvTranspose1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ConvTranspose1dImpl +:members: +:undoc-members: +``` + +## ConvTranspose2d + +Applies 2D transposed convolution for upsampling. Commonly used in decoder +networks and generative models. + +```{doxygenclass} torch::nn::ConvTranspose2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ConvTranspose2dImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +// Create ConvTranspose2d for upsampling +auto conv_transpose = torch::nn::ConvTranspose2d( + torch::nn::ConvTranspose2dOptions(64, 32, 4) + .stride(2) + .padding(1)); +``` + +## ConvTranspose3d + +Applies 3D transposed convolution for upsampling volumetric data. + +```{doxygenclass} torch::nn::ConvTranspose3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ConvTranspose3dImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/dropout.md b/docs/cpp/source/api/nn/dropout.md new file mode 100644 index 0000000000000..0777452ab9f34 --- /dev/null +++ b/docs/cpp/source/api/nn/dropout.md @@ -0,0 +1,86 @@ +--- +myst: + html_meta: + description: Dropout layers in PyTorch C++ — Dropout, Dropout2d, Dropout3d, AlphaDropout, and FeatureAlphaDropout. + keywords: PyTorch, C++, Dropout, Dropout2d, Dropout3d, AlphaDropout, regularization +--- + +# Dropout Layers + +Dropout randomly zeros elements during training as a regularization technique, +preventing overfitting by forcing the network to learn redundant representations. +During evaluation, dropout is disabled and outputs are scaled appropriately. + +- **Dropout**: Standard dropout for fully-connected layers +- **Dropout2d/3d**: Spatial dropout that zeros entire channels (better for CNNs) +- **AlphaDropout**: Maintains self-normalizing property (use with SELU activation) + +```{note} + +Remember to call `model->train()` during training and `model->eval()` during +inference to properly enable/disable dropout behavior. +``` + +## Dropout + +```{doxygenclass} torch::nn::Dropout +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::DropoutImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto dropout = torch::nn::Dropout(torch::nn::DropoutOptions(0.5)); +``` + +## Dropout2d / Dropout3d + +```{doxygenclass} torch::nn::Dropout2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::Dropout2dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::Dropout3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::Dropout3dImpl +:members: +:undoc-members: +``` + +## AlphaDropout + +```{doxygenclass} torch::nn::AlphaDropout +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AlphaDropoutImpl +:members: +:undoc-members: +``` + +## FeatureAlphaDropout + +```{doxygenclass} torch::nn::FeatureAlphaDropout +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::FeatureAlphaDropoutImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/embedding.md b/docs/cpp/source/api/nn/embedding.md new file mode 100644 index 0000000000000..f766a4ea17bfc --- /dev/null +++ b/docs/cpp/source/api/nn/embedding.md @@ -0,0 +1,55 @@ +--- +myst: + html_meta: + description: Embedding layers in PyTorch C++ — Embedding and EmbeddingBag for sparse and dense lookups. + keywords: PyTorch, C++, Embedding, EmbeddingBag, lookup table, sparse, NLP +--- + +# Embedding Layers + +Embedding layers map discrete tokens (words, categories, IDs) to dense vector +representations. They are the foundation of NLP models and recommendation systems. + +- **Embedding**: Standard lookup table that maps indices to dense vectors +- **EmbeddingBag**: Computes sums or means of embeddings (efficient for sparse features) + +**Key parameters:** + +- `num_embeddings`: Size of the vocabulary (number of unique tokens) +- `embedding_dim`: Dimension of each embedding vector +- `padding_idx`: Index that outputs zeros (useful for padding tokens) + +## Embedding + +```{doxygenclass} torch::nn::Embedding +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::EmbeddingImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto embedding = torch::nn::Embedding( + torch::nn::EmbeddingOptions(10000, 256) // num_embeddings, embedding_dim + .padding_idx(0)); + +auto indices = torch::tensor({1, 2, 3, 4}); +auto embedded = embedding->forward(indices); // [4, 256] +``` + +## EmbeddingBag + +```{doxygenclass} torch::nn::EmbeddingBag +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::EmbeddingBagImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/functional.md b/docs/cpp/source/api/nn/functional.md new file mode 100644 index 0000000000000..0ace995624041 --- /dev/null +++ b/docs/cpp/source/api/nn/functional.md @@ -0,0 +1,363 @@ +--- +myst: + html_meta: + description: Functional API in PyTorch C++ — torch::nn::functional stateless operations for neural networks. + keywords: PyTorch, C++, functional, torch::nn::functional, relu, conv2d, linear, softmax +--- + +# Functional API + +The `torch::nn::functional` namespace provides stateless versions of neural +network operations. Unlike module classes, functional operations do not hold +learnable parameters — you pass weights explicitly. + +**When to use functional vs modules:** + +- Use **modules** (`torch::nn::Conv2d`) when you need learnable parameters + managed automatically (training, saving, loading). +- Use **functional** (`torch::nn::functional::conv2d`) when you already have + weights as tensors, or for operations without parameters (e.g., `relu`). + +```cpp +#include +namespace F = torch::nn::functional; + +// Stateless activation — no module needed +auto output = F::relu(input); + +// Convolution with explicit weight tensor +auto output = F::conv2d(input, weight, F::Conv2dFuncOptions().stride(1).padding(1)); + +// Softmax along a dimension +auto probs = F::softmax(logits, F::SoftmaxFuncOptions(/*dim=*/1)); +``` + +## Activation Functions + +```{doxygenfunction} torch::nn::functional::elu +``` +```{doxygenfunction} torch::nn::functional::selu +``` +```{doxygenfunction} torch::nn::functional::hardshrink +``` +```{doxygenfunction} torch::nn::functional::hardtanh +``` +```{doxygenfunction} torch::nn::functional::leaky_relu +``` +```{doxygenfunction} torch::nn::functional::logsigmoid +``` +```{doxygenfunction} torch::nn::functional::glu +``` +```{doxygenfunction} torch::nn::functional::gelu +``` +```{doxygenfunction} torch::nn::functional::silu +``` +```{doxygenfunction} torch::nn::functional::mish +``` +```{doxygenfunction} torch::nn::functional::prelu +``` +```{doxygenfunction} torch::nn::functional::relu +``` +```{doxygenfunction} torch::nn::functional::relu6 +``` +```{doxygenfunction} torch::nn::functional::rrelu +``` +```{doxygenfunction} torch::nn::functional::celu +``` +```{doxygenfunction} torch::nn::functional::softplus +``` +```{doxygenfunction} torch::nn::functional::softshrink +``` +```{doxygenfunction} torch::nn::functional::softsign +``` +```{doxygenfunction} torch::nn::functional::tanhshrink +``` +```{doxygenfunction} torch::nn::functional::threshold +``` +```{doxygenfunction} torch::nn::functional::softmax +``` +```{doxygenfunction} torch::nn::functional::softmin +``` +```{doxygenfunction} torch::nn::functional::log_softmax +``` +```{doxygenfunction} torch::nn::functional::gumbel_softmax +``` + +## Convolution Functions + +```{doxygenfunction} torch::nn::functional::conv1d +``` +```{doxygenfunction} torch::nn::functional::conv2d +``` +```{doxygenfunction} torch::nn::functional::conv3d +``` +```{doxygenfunction} torch::nn::functional::conv_transpose1d +``` +```{doxygenfunction} torch::nn::functional::conv_transpose2d +``` +```{doxygenfunction} torch::nn::functional::conv_transpose3d +``` + +## Pooling Functions + +```{doxygenfunction} torch::nn::functional::avg_pool1d +``` +```{doxygenfunction} torch::nn::functional::avg_pool2d +``` +```{doxygenfunction} torch::nn::functional::avg_pool3d +``` +```{doxygenfunction} torch::nn::functional::max_pool1d +``` +```{doxygenfunction} torch::nn::functional::max_pool2d +``` +```{doxygenfunction} torch::nn::functional::max_pool3d +``` +```{doxygenfunction} torch::nn::functional::max_pool1d_with_indices +``` +```{doxygenfunction} torch::nn::functional::max_pool2d_with_indices +``` +```{doxygenfunction} torch::nn::functional::max_pool3d_with_indices +``` +```{doxygenfunction} torch::nn::functional::adaptive_max_pool1d +``` +```{doxygenfunction} torch::nn::functional::adaptive_max_pool2d +``` +```{doxygenfunction} torch::nn::functional::adaptive_max_pool3d +``` +```{doxygenfunction} torch::nn::functional::adaptive_avg_pool1d +``` +```{doxygenfunction} torch::nn::functional::adaptive_avg_pool2d +``` +```{doxygenfunction} torch::nn::functional::adaptive_avg_pool3d +``` +```{doxygenfunction} torch::nn::functional::max_unpool1d +``` +```{doxygenfunction} torch::nn::functional::max_unpool2d +``` +```{doxygenfunction} torch::nn::functional::max_unpool3d +``` +```{doxygenfunction} torch::nn::functional::fractional_max_pool2d +``` +```{doxygenfunction} torch::nn::functional::fractional_max_pool3d +``` +```{doxygenfunction} torch::nn::functional::lp_pool1d +``` +```{doxygenfunction} torch::nn::functional::lp_pool2d +``` +```{doxygenfunction} torch::nn::functional::lp_pool3d +``` + +## Linear Functions + +```{doxygenfunction} torch::nn::functional::linear +``` +```{doxygenfunction} torch::nn::functional::bilinear +``` + +## Dropout Functions + +```{doxygenfunction} torch::nn::functional::dropout +``` +```{doxygenfunction} torch::nn::functional::dropout2d +``` +```{doxygenfunction} torch::nn::functional::dropout3d +``` +```{doxygenfunction} torch::nn::functional::alpha_dropout +``` +```{doxygenfunction} torch::nn::functional::feature_alpha_dropout +``` + +## Embedding Functions + +```{doxygenfunction} torch::nn::functional::one_hot +``` +```{doxygenfunction} torch::nn::functional::embedding +``` +```{doxygenfunction} torch::nn::functional::embedding_bag +``` + +## Normalization Functions + +```{doxygenfunction} torch::nn::functional::batch_norm +``` +```{doxygenfunction} torch::nn::functional::instance_norm +``` +```{doxygenfunction} torch::nn::functional::layer_norm +``` +```{doxygenfunction} torch::nn::functional::group_norm +``` +```{doxygenfunction} torch::nn::functional::local_response_norm +``` +```{doxygenfunction} torch::nn::functional::normalize +``` + +## Loss Functions + +```{doxygenfunction} torch::nn::functional::l1_loss +``` +```{doxygenfunction} torch::nn::functional::mse_loss +``` +```{doxygenfunction} torch::nn::functional::binary_cross_entropy +``` +```{doxygenfunction} torch::nn::functional::binary_cross_entropy_with_logits +``` +```{doxygenfunction} torch::nn::functional::cross_entropy +``` +```{doxygenfunction} torch::nn::functional::nll_loss +``` +```{doxygenfunction} torch::nn::functional::kl_div +``` +```{doxygenfunction} torch::nn::functional::smooth_l1_loss(const Tensor& input, const Tensor& target, const SmoothL1LossFuncOptions& options) +``` +```{doxygenfunction} torch::nn::functional::huber_loss +``` +```{doxygenfunction} torch::nn::functional::hinge_embedding_loss +``` +```{doxygenfunction} torch::nn::functional::multi_margin_loss +``` +```{doxygenfunction} torch::nn::functional::cosine_embedding_loss +``` +```{doxygenfunction} torch::nn::functional::margin_ranking_loss +``` +```{doxygenfunction} torch::nn::functional::multilabel_margin_loss +``` +```{doxygenfunction} torch::nn::functional::soft_margin_loss +``` +```{doxygenfunction} torch::nn::functional::multilabel_soft_margin_loss +``` +```{doxygenfunction} torch::nn::functional::triplet_margin_loss +``` +```{doxygenfunction} torch::nn::functional::triplet_margin_with_distance_loss +``` +```{doxygenfunction} torch::nn::functional::ctc_loss +``` +```{doxygenfunction} torch::nn::functional::poisson_nll_loss +``` + +## Distance Functions + +```{doxygenfunction} torch::nn::functional::cosine_similarity +``` +```{doxygenfunction} torch::nn::functional::pairwise_distance +``` +```{doxygenfunction} torch::nn::functional::pdist +``` + +## Vision Functions + +```{doxygenfunction} torch::nn::functional::interpolate +``` +```{doxygenfunction} torch::nn::functional::affine_grid +``` +```{doxygenfunction} torch::nn::functional::grid_sample +``` +```{doxygenfunction} torch::nn::functional::pad +``` +```{doxygenfunction} torch::nn::functional::pixel_shuffle +``` +```{doxygenfunction} torch::nn::functional::pixel_unshuffle +``` + +## Fold/Unfold + +```{doxygenfunction} torch::nn::functional::fold +``` +```{doxygenfunction} torch::nn::functional::unfold +``` + +## Functional Options Structs + +Each functional operation that takes configuration uses a corresponding options +struct. The naming convention is `FuncOptions`. + +**Activation Options:** + +```{doxygentypedef} torch::nn::functional::ELUFuncOptions +``` +```{doxygentypedef} torch::nn::functional::SELUFuncOptions +``` +```{doxygentypedef} torch::nn::functional::GLUFuncOptions +``` +```{doxygentypedef} torch::nn::functional::GELUFuncOptions +``` +```{doxygentypedef} torch::nn::functional::HardshrinkFuncOptions +``` +```{doxygentypedef} torch::nn::functional::HardtanhFuncOptions +``` +```{doxygentypedef} torch::nn::functional::LeakyReLUFuncOptions +``` +```{doxygentypedef} torch::nn::functional::ReLUFuncOptions +``` +```{doxygentypedef} torch::nn::functional::ReLU6FuncOptions +``` +```{doxygentypedef} torch::nn::functional::CELUFuncOptions +``` +```{doxygentypedef} torch::nn::functional::SoftplusFuncOptions +``` +```{doxygentypedef} torch::nn::functional::SoftshrinkFuncOptions +``` +```{doxygentypedef} torch::nn::functional::ThresholdFuncOptions +``` + +**Convolution Options:** + +```{doxygentypedef} torch::nn::functional::Conv1dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::Conv2dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::Conv3dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::ConvTranspose1dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::ConvTranspose2dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::ConvTranspose3dFuncOptions +``` + +**Pooling Options:** + +```{doxygentypedef} torch::nn::functional::AvgPool1dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::AvgPool2dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::AvgPool3dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::MaxPool1dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::MaxPool2dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::MaxPool3dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::AdaptiveMaxPool1dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::AdaptiveMaxPool2dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::AdaptiveMaxPool3dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::AdaptiveAvgPool1dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::AdaptiveAvgPool2dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::AdaptiveAvgPool3dFuncOptions +``` + +**Other Options:** + +```{doxygentypedef} torch::nn::functional::CosineSimilarityFuncOptions +``` +```{doxygentypedef} torch::nn::functional::PairwiseDistanceFuncOptions +``` +```{doxygentypedef} torch::nn::functional::Dropout2dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::Dropout3dFuncOptions +``` +```{doxygentypedef} torch::nn::functional::L1LossFuncOptions +``` +```{doxygentypedef} torch::nn::functional::FoldFuncOptions +``` +```{doxygentypedef} torch::nn::functional::UnfoldFuncOptions +``` +```{doxygentypedef} torch::nn::functional::PixelShuffleFuncOptions +``` +```{doxygentypedef} torch::nn::functional::PixelUnshuffleFuncOptions +``` diff --git a/docs/cpp/source/api/nn/index.md b/docs/cpp/source/api/nn/index.md new file mode 100644 index 0000000000000..c5bf65125e12b --- /dev/null +++ b/docs/cpp/source/api/nn/index.md @@ -0,0 +1,94 @@ +--- +myst: + html_meta: + description: PyTorch C++ neural network modules — torch::nn API for defining and training models. + keywords: PyTorch, C++, nn, Module, neural network, torch::nn +--- + +# Neural Network Modules (torch::nn) + +The `torch::nn` namespace provides neural network building blocks that mirror +Python's `torch.nn` module. It uses a PIMPL (Pointer to Implementation) pattern +where user-facing classes like `Conv2d` wrap internal `Conv2dImpl` classes. + +**When to use torch::nn:** + +- Building neural network models in C++ +- Creating custom layers and modules +- Porting Python models to C++ for production inference +- Training models entirely in C++ + +**Basic usage:** + +```cpp +#include + +// Define a simple model +struct Net : torch::nn::Module { + torch::nn::Conv2d conv1{nullptr}; + torch::nn::Linear fc1{nullptr}; + + Net() { + conv1 = register_module("conv1", torch::nn::Conv2d( + torch::nn::Conv2dOptions(1, 32, 3).stride(1).padding(1))); + fc1 = register_module("fc1", torch::nn::Linear(32 * 28 * 28, 10)); + } + + torch::Tensor forward(torch::Tensor x) { + x = torch::relu(conv1->forward(x)); + x = x.view({-1, 32 * 28 * 28}); + return fc1->forward(x); + } +}; + +// Create and use the model +auto model = std::make_shared(); +auto input = torch::randn({1, 1, 28, 28}); +auto output = model->forward(input); +``` + +## Header Files + +- `torch/nn.h` - Main neural network header (includes all modules) +- `torch/nn/module.h` - Base Module class +- `torch/nn/modules.h` - All module implementations +- `torch/nn/options.h` - Options structs for modules +- `torch/nn/functional.h` - Functional API + +## Module Base Class + +All neural network modules inherit from `torch::nn::Module`, which provides +parameter management, serialization, device/dtype conversion, and hooks. + +```{doxygenclass} torch::nn::Module +``` + +**Key features:** + +- `register_module()`: Register submodules for parameter tracking +- `register_parameter()`: Register learnable parameters +- `register_buffer()`: Register non-learnable state (e.g., running mean) +- `parameters()` / `named_parameters()`: Iterate over all parameters +- `to()`: Move module to a device or convert dtype +- `train()` / `eval()`: Toggle training/evaluation mode +- `save()` / `load()`: Serialize and deserialize module state + +## Module Categories + +```{toctree} +:maxdepth: 1 + +containers +convolution +pooling +linear +activation +normalization +dropout +embedding +recurrent +transformer +loss +functional +utilities +``` diff --git a/docs/cpp/source/api/nn/linear.md b/docs/cpp/source/api/nn/linear.md new file mode 100644 index 0000000000000..8aa0c55a0ba4a --- /dev/null +++ b/docs/cpp/source/api/nn/linear.md @@ -0,0 +1,84 @@ +--- +myst: + html_meta: + description: Linear layers in PyTorch C++ — Linear, Bilinear, and Flatten modules. + keywords: PyTorch, C++, Linear, Bilinear, Flatten, fully connected, dense layer +--- + +# Linear Layers + +Linear layers apply affine transformations to input data: `y = xW^T + b`. +They are the building blocks of fully-connected networks and are used for +feature transformation, classification heads, and projection layers. + +- **Linear**: Standard fully-connected layer +- **Bilinear**: Bilinear transformation of two inputs +- **Identity**: Pass-through layer (useful for skip connections) +- **Flatten/Unflatten**: Reshape tensors between convolutional and linear layers + +## Linear + +```{doxygenclass} torch::nn::Linear +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LinearImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto linear = torch::nn::Linear(torch::nn::LinearOptions(784, 256).bias(true)); +auto output = linear->forward(input); // input: [N, 784] +``` + +## Bilinear + +```{doxygenclass} torch::nn::Bilinear +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::BilinearImpl +:members: +:undoc-members: +``` + +## Identity + +```{doxygenclass} torch::nn::Identity +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::IdentityImpl +:members: +:undoc-members: +``` + +## Flatten + +```{doxygenclass} torch::nn::Flatten +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::FlattenImpl +:members: +:undoc-members: +``` + +## Unflatten + +```{doxygenclass} torch::nn::Unflatten +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::UnflattenImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/loss.md b/docs/cpp/source/api/nn/loss.md new file mode 100644 index 0000000000000..8f1ab82bc234d --- /dev/null +++ b/docs/cpp/source/api/nn/loss.md @@ -0,0 +1,284 @@ +--- +myst: + html_meta: + description: Loss functions in PyTorch C++ — CrossEntropyLoss, MSELoss, NLLLoss, BCELoss, and more. + keywords: PyTorch, C++, loss, CrossEntropyLoss, MSELoss, NLLLoss, BCELoss, L1Loss +--- + +# Loss Functions + +Loss functions measure how well the model's predictions match the targets. +The choice of loss function depends on your task type and data characteristics. + +**Regression losses:** + +- **L1Loss/MSELoss**: Basic regression losses (MAE vs MSE) +- **SmoothL1Loss/HuberLoss**: Robust to outliers + +**Classification losses:** + +- **CrossEntropyLoss**: Multi-class classification (combines LogSoftmax + NLLLoss) +- **NLLLoss**: Negative log likelihood (use with LogSoftmax output) +- **BCELoss/BCEWithLogitsLoss**: Binary classification + +**Specialized losses:** + +- **CTCLoss**: Sequence-to-sequence without alignment (speech recognition) +- **TripletMarginLoss**: Metric learning (similarity/embedding tasks) +- **CosineEmbeddingLoss**: Similarity learning with cosine distance + +## L1Loss + +```{doxygenclass} torch::nn::L1Loss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::L1LossImpl +:members: +:undoc-members: +``` + +## MSELoss + +```{doxygenclass} torch::nn::MSELoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::MSELossImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto loss_fn = torch::nn::MSELoss(); +auto loss = loss_fn->forward(predictions, targets); +``` + +## CrossEntropyLoss + +```{doxygenclass} torch::nn::CrossEntropyLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::CrossEntropyLossImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto loss_fn = torch::nn::CrossEntropyLoss(); +auto logits = torch::randn({32, 10}); // [batch, num_classes] +auto targets = torch::randint(0, 10, {32}); // [batch] +auto loss = loss_fn->forward(logits, targets); +``` + +## NLLLoss + +```{doxygenclass} torch::nn::NLLLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::NLLLossImpl +:members: +:undoc-members: +``` + +## BCELoss + +```{doxygenclass} torch::nn::BCELoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::BCELossImpl +:members: +:undoc-members: +``` + +## BCEWithLogitsLoss + +```{doxygenclass} torch::nn::BCEWithLogitsLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::BCEWithLogitsLossImpl +:members: +:undoc-members: +``` + +## HuberLoss + +```{doxygenclass} torch::nn::HuberLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::HuberLossImpl +:members: +:undoc-members: +``` + +## SmoothL1Loss + +```{doxygenclass} torch::nn::SmoothL1Loss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::SmoothL1LossImpl +:members: +:undoc-members: +``` + +## KLDivLoss + +```{doxygenclass} torch::nn::KLDivLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::KLDivLossImpl +:members: +:undoc-members: +``` + +## CTCLoss + +```{doxygenclass} torch::nn::CTCLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::CTCLossImpl +:members: +:undoc-members: +``` + +## PoissonNLLLoss + +```{doxygenclass} torch::nn::PoissonNLLLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::PoissonNLLLossImpl +:members: +:undoc-members: +``` + +## MarginRankingLoss + +```{doxygenclass} torch::nn::MarginRankingLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::MarginRankingLossImpl +:members: +:undoc-members: +``` + +## HingeEmbeddingLoss + +```{doxygenclass} torch::nn::HingeEmbeddingLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::HingeEmbeddingLossImpl +:members: +:undoc-members: +``` + +## CosineEmbeddingLoss + +```{doxygenclass} torch::nn::CosineEmbeddingLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::CosineEmbeddingLossImpl +:members: +:undoc-members: +``` + +## MultiMarginLoss + +```{doxygenclass} torch::nn::MultiMarginLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::MultiMarginLossImpl +:members: +:undoc-members: +``` + +## MultiLabelMarginLoss + +```{doxygenclass} torch::nn::MultiLabelMarginLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::MultiLabelMarginLossImpl +:members: +:undoc-members: +``` + +## MultiLabelSoftMarginLoss + +```{doxygenclass} torch::nn::MultiLabelSoftMarginLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::MultiLabelSoftMarginLossImpl +:members: +:undoc-members: +``` + +## SoftMarginLoss + +```{doxygenclass} torch::nn::SoftMarginLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::SoftMarginLossImpl +:members: +:undoc-members: +``` + +## TripletMarginLoss + +```{doxygenclass} torch::nn::TripletMarginLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::TripletMarginLossImpl +:members: +:undoc-members: +``` + +## TripletMarginWithDistanceLoss + +```{doxygenclass} torch::nn::TripletMarginWithDistanceLoss +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::TripletMarginWithDistanceLossImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/normalization.md b/docs/cpp/source/api/nn/normalization.md new file mode 100644 index 0000000000000..3ccfe99957e06 --- /dev/null +++ b/docs/cpp/source/api/nn/normalization.md @@ -0,0 +1,142 @@ +--- +myst: + html_meta: + description: Normalization layers in PyTorch C++ — BatchNorm, LayerNorm, GroupNorm, InstanceNorm, and LocalResponseNorm. + keywords: PyTorch, C++, normalization, BatchNorm, LayerNorm, GroupNorm, InstanceNorm +--- + +# Normalization Layers + +Normalization layers stabilize and accelerate training by normalizing intermediate +activations. They help with gradient flow and allow higher learning rates. + +- **BatchNorm**: Normalizes across batch dimension; most common in CNNs +- **InstanceNorm**: Normalizes each sample independently; popular in style transfer +- **LayerNorm**: Normalizes across feature dimension; standard in transformers +- **GroupNorm**: Normalizes within groups of channels; works with small batches +- **LocalResponseNorm**: Lateral inhibition inspired by neuroscience (less common today) + +## BatchNorm1d / BatchNorm2d / BatchNorm3d + +```{doxygenclass} torch::nn::BatchNorm1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::BatchNorm1dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::BatchNorm2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::BatchNorm2dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::BatchNorm3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::BatchNorm3dImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto bn = torch::nn::BatchNorm2d( + torch::nn::BatchNorm2dOptions(64) // num_features + .eps(1e-5) + .momentum(0.1) + .affine(true) + .track_running_stats(true)); +``` + +## InstanceNorm1d / InstanceNorm2d / InstanceNorm3d + +```{doxygenclass} torch::nn::InstanceNorm1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::InstanceNorm1dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::InstanceNorm2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::InstanceNorm2dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::InstanceNorm3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::InstanceNorm3dImpl +:members: +:undoc-members: +``` + +## LayerNorm + +```{doxygenclass} torch::nn::LayerNorm +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LayerNormImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto ln = torch::nn::LayerNorm( + torch::nn::LayerNormOptions({768})); // normalized_shape +``` + +## GroupNorm + +```{doxygenclass} torch::nn::GroupNorm +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::GroupNormImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto gn = torch::nn::GroupNorm( + torch::nn::GroupNormOptions(32, 256)); // num_groups, num_channels +``` + +## LocalResponseNorm + +```{doxygenclass} torch::nn::LocalResponseNorm +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LocalResponseNormImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/pooling.md b/docs/cpp/source/api/nn/pooling.md new file mode 100644 index 0000000000000..679b24833d7ab --- /dev/null +++ b/docs/cpp/source/api/nn/pooling.md @@ -0,0 +1,262 @@ +--- +myst: + html_meta: + description: Pooling layers in PyTorch C++ — MaxPool, AvgPool, AdaptiveMaxPool, AdaptiveAvgPool, and LPPool. + keywords: PyTorch, C++, pooling, MaxPool2d, AvgPool2d, AdaptiveAvgPool2d, max pooling +--- + +# Pooling Layers + +Pooling layers reduce spatial dimensions by aggregating values in local regions, +providing translation invariance and reducing computational cost in deeper layers. + +- **MaxPool**: Takes the maximum value in each pooling window (preserves strong features) +- **AvgPool**: Takes the average value in each pooling window (smoother downsampling) +- **AdaptivePool**: Automatically calculates kernel size to produce a target output size +- **FractionalMaxPool**: Randomized pooling with fractional output size +- **MaxUnpool**: Computes the partial inverse of MaxPool using stored indices +- **LPPool**: Power-average pooling (generalization of avg/max pooling) + +## MaxPool1d / MaxPool2d / MaxPool3d + +```{doxygenclass} torch::nn::MaxPool1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MaxPool1dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MaxPool2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MaxPool2dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MaxPool3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MaxPool3dImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto pool = torch::nn::MaxPool2d( + torch::nn::MaxPool2dOptions(2).stride(2)); +``` + +## AvgPool1d / AvgPool2d / AvgPool3d + +```{doxygenclass} torch::nn::AvgPool1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AvgPool1dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AvgPool2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AvgPool2dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AvgPool3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AvgPool3dImpl +:members: +:undoc-members: +``` + +## AdaptiveAvgPool1d / AdaptiveAvgPool2d / AdaptiveAvgPool3d + +```{doxygenclass} torch::nn::AdaptiveAvgPool1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AdaptiveAvgPool1dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AdaptiveAvgPool2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AdaptiveAvgPool2dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AdaptiveAvgPool3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AdaptiveAvgPool3dImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +// Output will always be 7x7 regardless of input size +auto adaptive_pool = torch::nn::AdaptiveAvgPool2d( + torch::nn::AdaptiveAvgPool2dOptions({7, 7})); +``` + +## AdaptiveMaxPool1d / AdaptiveMaxPool2d / AdaptiveMaxPool3d + +```{doxygenclass} torch::nn::AdaptiveMaxPool1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AdaptiveMaxPool1dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AdaptiveMaxPool2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AdaptiveMaxPool2dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AdaptiveMaxPool3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::AdaptiveMaxPool3dImpl +:members: +:undoc-members: +``` + +## FractionalMaxPool2d / FractionalMaxPool3d + +```{doxygenclass} torch::nn::FractionalMaxPool2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::FractionalMaxPool2dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::FractionalMaxPool3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::FractionalMaxPool3dImpl +:members: +:undoc-members: +``` + +## MaxUnpool1d / MaxUnpool2d / MaxUnpool3d + +Computes a partial inverse of `MaxPool`, using the indices of the maximum +values computed during pooling to place values back into unpooled positions. + +```{doxygenclass} torch::nn::MaxUnpool1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MaxUnpool1dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MaxUnpool2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MaxUnpool2dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MaxUnpool3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MaxUnpool3dImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto pool = torch::nn::MaxPool2d( + torch::nn::MaxPool2dOptions(2).stride(2).return_indices(true)); +auto unpool = torch::nn::MaxUnpool2d( + torch::nn::MaxUnpoolOptions<2>(2).stride(2)); + +auto [output, indices] = pool->forward_with_indices(input); +auto reconstructed = unpool->forward(output, indices); +``` + +## LPPool1d / LPPool2d / LPPool3d + +```{doxygenclass} torch::nn::LPPool1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LPPool1dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LPPool2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LPPool2dImpl +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LPPool3d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LPPool3dImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/recurrent.md b/docs/cpp/source/api/nn/recurrent.md new file mode 100644 index 0000000000000..f286ebb6b52b5 --- /dev/null +++ b/docs/cpp/source/api/nn/recurrent.md @@ -0,0 +1,125 @@ +--- +myst: + html_meta: + description: Recurrent layers in PyTorch C++ — RNN, LSTM, and GRU modules for sequence modeling. + keywords: PyTorch, C++, RNN, LSTM, GRU, recurrent, sequence, hidden state +--- + +# Recurrent Layers + +Recurrent layers process sequential data by maintaining hidden state across time steps. +They are essential for tasks involving sequences: language modeling, speech recognition, +time series prediction, and more. + +- **RNN**: Basic recurrent layer (simple but prone to vanishing gradients) +- **LSTM**: Long Short-Term Memory (gated architecture, handles long-range dependencies) +- **GRU**: Gated Recurrent Unit (simpler than LSTM, often similar performance) +- **Cell variants**: Single-step versions for custom loop implementations + +**Key parameters:** + +- `input_size`: Number of features in input +- `hidden_size`: Number of features in hidden state +- `num_layers`: Number of stacked recurrent layers +- `batch_first`: If true, input shape is `[batch, seq, features]` +- `bidirectional`: Process sequence in both directions + +## RNN + +```{doxygenclass} torch::nn::RNN +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::RNNImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto rnn = torch::nn::RNN( + torch::nn::RNNOptions(128, 256) // input_size, hidden_size + .num_layers(2) + .batch_first(true) + .bidirectional(false)); + +auto input = torch::randn({32, 10, 128}); // [batch, seq_len, input_size] +auto [output, hidden] = rnn->forward(input); +``` + +## LSTM + +```{doxygenclass} torch::nn::LSTM +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LSTMImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto lstm = torch::nn::LSTM( + torch::nn::LSTMOptions(128, 256) + .num_layers(2) + .batch_first(true) + .dropout(0.1) + .bidirectional(true)); + +auto input = torch::randn({32, 10, 128}); +auto [output, state] = lstm->forward(input); +auto [h_n, c_n] = state; // hidden state, cell state +``` + +## GRU + +```{doxygenclass} torch::nn::GRU +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::GRUImpl +:members: +:undoc-members: +``` + +## RNNCell + +```{doxygenclass} torch::nn::RNNCell +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::RNNCellImpl +:members: +:undoc-members: +``` + +## LSTMCell + +```{doxygenclass} torch::nn::LSTMCell +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::LSTMCellImpl +:members: +:undoc-members: +``` + +## GRUCell + +```{doxygenclass} torch::nn::GRUCell +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::GRUCellImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/transformer.md b/docs/cpp/source/api/nn/transformer.md new file mode 100644 index 0000000000000..664564d4b3519 --- /dev/null +++ b/docs/cpp/source/api/nn/transformer.md @@ -0,0 +1,112 @@ +--- +myst: + html_meta: + description: Transformer layers in PyTorch C++ — Transformer, TransformerEncoder, TransformerDecoder, and MultiheadAttention. + keywords: PyTorch, C++, Transformer, TransformerEncoder, TransformerDecoder, MultiheadAttention, attention +--- + +# Transformer Layers + +Transformer layers use self-attention mechanisms to process sequences in parallel, +enabling efficient training on long sequences. They are the foundation of modern +NLP models (BERT, GPT) and increasingly used in vision and other domains. + +- **Transformer**: Complete encoder-decoder architecture +- **TransformerEncoder/Decoder**: Standalone encoder or decoder stacks +- **TransformerEncoderLayer/DecoderLayer**: Individual transformer blocks +- **MultiheadAttention**: Core attention mechanism used throughout + +**Key parameters:** + +- `d_model`: Dimension of the model (embedding dimension) +- `nhead`: Number of attention heads +- `num_encoder_layers/num_decoder_layers`: Number of stacked layers +- `dim_feedforward`: Dimension of feedforward network +- `dropout`: Dropout rate for regularization + +## Transformer + +Complete encoder-decoder transformer architecture. + +```{doxygenclass} torch::nn::Transformer +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::TransformerImpl +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto transformer = torch::nn::Transformer( + torch::nn::TransformerOptions() + .d_model(512) + .nhead(8) + .num_encoder_layers(6) + .num_decoder_layers(6) + .dim_feedforward(2048) + .dropout(0.1)); +``` + +## TransformerEncoder + +Stack of encoder layers for processing source sequences. + +```{doxygenclass} torch::nn::TransformerEncoder +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::TransformerEncoderImpl +:members: +:undoc-members: +``` + +## TransformerDecoder + +Stack of decoder layers for generating target sequences. + +```{doxygenclass} torch::nn::TransformerDecoder +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::TransformerDecoderImpl +:members: +:undoc-members: +``` + +## TransformerEncoderLayer + +Single encoder layer with self-attention and feedforward network. + +```{doxygenclass} torch::nn::TransformerEncoderLayerImpl +:members: +:undoc-members: +``` + +## TransformerDecoderLayer + +Single decoder layer with self-attention, cross-attention, and feedforward network. + +```{doxygenclass} TransformerDecoderLayerImpl +:members: +:undoc-members: +``` + +## MultiheadAttention + +Scaled dot-product attention with multiple parallel heads. + +```{doxygenclass} torch::nn::MultiheadAttention +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::MultiheadAttentionImpl +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/nn/utilities.md b/docs/cpp/source/api/nn/utilities.md new file mode 100644 index 0000000000000..6bde6de870a0d --- /dev/null +++ b/docs/cpp/source/api/nn/utilities.md @@ -0,0 +1,261 @@ +--- +myst: + html_meta: + description: Neural network utilities in PyTorch C++ — parameter initialization, module cloning, padding layers, and vision utilities. + keywords: PyTorch, C++, nn utilities, init, Cloneable, AnyModule, padding, PixelShuffle, Upsample +--- + +# Utilities + +Additional utilities for building neural networks: parameter initialization, +module cloning, type-erased containers, padding layers, and vision utilities. + +## Parameter Initialization + +The `torch::nn::init` namespace provides functions for initializing module parameters: + +```cpp +#include + +// Xavier/Glorot initialization +torch::nn::init::xavier_uniform_(linear->weight); +torch::nn::init::xavier_normal_(linear->weight); + +// Kaiming/He initialization +torch::nn::init::kaiming_uniform_(conv->weight, /*a=*/0, torch::kFanIn, torch::kReLU); +torch::nn::init::kaiming_normal_(conv->weight); + +// Other initializations +torch::nn::init::zeros_(linear->bias); +torch::nn::init::ones_(bn->weight); +torch::nn::init::constant_(linear->bias, 0.1); +torch::nn::init::normal_(linear->weight, /*mean=*/0, /*std=*/0.01); +torch::nn::init::uniform_(linear->weight, /*a=*/-0.1, /*b=*/0.1); +torch::nn::init::orthogonal_(rnn->weight_hh); +``` + +## Cloneable + +```{doxygenclass} torch::nn::Cloneable +``` + +All `torch::nn` modules inherit from `Cloneable`, enabling deep copies: + +```cpp +auto model = torch::nn::Linear(10, 5); +auto model_copy = std::dynamic_pointer_cast(model->clone()); +``` + +## AnyModule + +`AnyModule` provides type-erased storage for any module, allowing you to +store heterogeneous modules in a single container. + +```{doxygenclass} torch::nn::AnyModule +:members: +:undoc-members: +``` + +**Example:** + +```cpp +torch::nn::AnyModule any_module(torch::nn::Linear(10, 5)); +auto output = any_module.forward(input); +``` + +## Functional + +Wraps a function or callable as a module, useful for inserting arbitrary +functions into a `Sequential` container. + +```{doxygenclass} torch::nn::FunctionalImpl +:members: +:undoc-members: +``` + +## ModuleHolder + +```{doxygenclass} torch::nn::ModuleHolder +:members: +:undoc-members: +``` + +## CosineSimilarity + +```{doxygenclass} torch::nn::CosineSimilarity +:members: +:undoc-members: +``` + +## PairwiseDistance + +```{doxygenclass} torch::nn::PairwiseDistance +:members: +:undoc-members: +``` + +## PackedSequence + +```{cpp:class} torch::nn::utils::rnn::PackedSequence + +Holds the data and list of `batch_sizes` of a packed sequence. +All RNN modules accept packed sequences as inputs. +``` + +```{cpp:function} const Tensor& data() const + +Returns the packed tensor containing all sequence elements. +``` + +```{cpp:function} const Tensor& batch_sizes() const + +Returns a 1D tensor of batch sizes at each time step. +``` + +```{cpp:function} const Tensor& sorted_indices() const + +Returns indices used to sort sequences by length (descending). +``` + +```{cpp:function} const Tensor& unsorted_indices() const + +Returns indices to restore the original sequence order. +``` + +```{cpp:function} PackedSequence to(torch::Device device) const + +Moves the packed sequence to the specified device. +``` + +See also: `torch::nn::utils::rnn::pack_padded_sequence` and +`torch::nn::utils::rnn::pad_packed_sequence`. + +## Padding Layers + +### ReflectionPad1d / ReflectionPad2d / ReflectionPad3d + +```{doxygenclass} torch::nn::ReflectionPad1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ReflectionPad2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ReflectionPad3d +:members: +:undoc-members: +``` + +### ReplicationPad1d / ReplicationPad2d / ReplicationPad3d + +```{doxygenclass} torch::nn::ReplicationPad1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ReplicationPad2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ReplicationPad3d +:members: +:undoc-members: +``` + +### ZeroPad1d / ZeroPad2d / ZeroPad3d + +```{doxygenclass} torch::nn::ZeroPad1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ZeroPad2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ZeroPad3d +:members: +:undoc-members: +``` + +### ConstantPad1d / ConstantPad2d / ConstantPad3d + +```{doxygenclass} torch::nn::ConstantPad1d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ConstantPad2d +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::ConstantPad3d +:members: +:undoc-members: +``` + +## Vision Layers + +### PixelShuffle + +```{doxygenclass} torch::nn::PixelShuffle +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::PixelShuffleOptions +:members: +:undoc-members: +``` + +### PixelUnshuffle + +```{doxygenclass} torch::nn::PixelUnshuffle +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::PixelUnshuffleOptions +:members: +:undoc-members: +``` + +### Upsample + +```{doxygenclass} torch::nn::Upsample +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::UpsampleOptions +:members: +:undoc-members: +``` + +### Fold / Unfold + +```{doxygenclass} torch::nn::Fold +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::FoldOptions +:members: +:undoc-members: +``` + +```{doxygenclass} torch::nn::Unfold +:members: +:undoc-members: +``` + +```{doxygenstruct} torch::nn::UnfoldOptions +:members: +:undoc-members: +``` diff --git a/docs/cpp/source/api/optim/adaptive.md b/docs/cpp/source/api/optim/adaptive.md new file mode 100644 index 0000000000000..3f48534567627 --- /dev/null +++ b/docs/cpp/source/api/optim/adaptive.md @@ -0,0 +1,158 @@ +--- +myst: + html_meta: + description: Adaptive optimizers in PyTorch C++ — Adam, AdamW, Adagrad, and RMSprop. + keywords: PyTorch, C++, Adam, AdamW, Adagrad, RMSprop, adaptive optimizer +--- + +# Adaptive Learning Rate Optimizers + +These optimizers automatically adapt the learning rate for each parameter based +on historical gradient information. They typically require less hyperparameter +tuning and work well across a wide range of problems. + +## Adam (Adaptive Moment Estimation) + +Adam combines the benefits of RMSprop and momentum, maintaining per-parameter +adaptive learning rates. It's an excellent default choice for most deep learning +tasks, especially when you want fast convergence with minimal tuning. + +**When to use:** + +- Transformers and attention-based models +- Quick prototyping and experimentation +- When you don't have time for extensive hyperparameter search +- General-purpose deep learning + +**Key parameters:** + +- `lr`: Learning rate (typical: 1e-3 to 1e-4) +- `betas`: Coefficients for running averages (default: {0.9, 0.999}) +- `eps`: Numerical stability term (default: 1e-8) +- `weight_decay`: L2 regularization (note: applied differently than in SGD) + +```{doxygenclass} torch::optim::Adam +:members: +:undoc-members: +``` + +**Example:** + +```cpp +// Standard Adam configuration +auto optimizer = torch::optim::Adam( + model->parameters(), + torch::optim::AdamOptions(1e-3) // learning rate + .betas({0.9, 0.999}) // momentum terms + .eps(1e-8) // numerical stability + .weight_decay(0)); // L2 penalty + +// For transformers, lower learning rate with warmup +auto optimizer = torch::optim::Adam( + model->parameters(), + torch::optim::AdamOptions(1e-4) + .betas({0.9, 0.98})); // β2=0.98 for transformers +``` + +## AdamW (Adam with Decoupled Weight Decay) + +AdamW fixes a subtle issue with Adam's weight decay implementation. In Adam, +weight decay is coupled with the gradient update, which can lead to suboptimal +regularization. AdamW decouples weight decay, applying it directly to the +weights as in SGD. + +**When to use:** + +- Always prefer AdamW over Adam when using weight decay +- Training transformers (BERT, GPT, etc.) +- When you want proper L2 regularization behavior + +**Key difference from Adam:** + +- In Adam: `weight = weight - lr * (grad + weight_decay * weight)` +- In AdamW: `weight = weight - lr * grad - lr * weight_decay * weight` + +```{doxygenclass} torch::optim::AdamW +:members: +:undoc-members: +``` + +**Example:** + +```cpp +// AdamW with decoupled weight decay - preferred for transformers +auto optimizer = torch::optim::AdamW( + model->parameters(), + torch::optim::AdamWOptions(1e-4) + .betas({0.9, 0.999}) + .weight_decay(0.01)); // Decoupled regularization +``` + +## RMSprop (Root Mean Square Propagation) + +RMSprop adapts the learning rate by dividing by a running average of recent +gradient magnitudes. It's particularly effective for recurrent neural networks +and problems with non-stationary objectives. + +**When to use:** + +- Training RNNs and LSTMs +- Non-stationary problems where gradient scale varies significantly +- Online learning scenarios + +**Key parameters:** + +- `lr`: Learning rate (typical: 1e-3 to 1e-2) +- `alpha`: Smoothing constant (default: 0.99) +- `momentum`: Optional momentum term +- `centered`: Use centered RMSprop (normalizes by variance) + +```{doxygenclass} torch::optim::RMSprop +:members: +:undoc-members: +``` + +**Example:** + +```cpp +// RMSprop for RNN training +auto optimizer = torch::optim::RMSprop( + model->parameters(), + torch::optim::RMSpropOptions(1e-3) + .alpha(0.99) // smoothing constant + .momentum(0.9) // optional momentum + .centered(true)); // normalize by variance +``` + +## Adagrad (Adaptive Gradient) + +Adagrad adapts the learning rate based on the accumulated sum of squared +gradients. Parameters with frequent updates get smaller learning rates, while +parameters with infrequent updates get larger rates. This makes it ideal for +sparse data. + +**When to use:** + +- NLP tasks with sparse features +- Embedding layers with infrequent updates +- Recommendation systems with sparse user/item features + +**Limitation:** Learning rate monotonically decreases, which can cause training +to stop prematurely. For long training runs, consider Adam or RMSprop. + +```{doxygenclass} torch::optim::Adagrad +:members: +:undoc-members: +``` + +**Example:** + +```cpp +// Adagrad for sparse NLP features +auto optimizer = torch::optim::Adagrad( + model->parameters(), + torch::optim::AdagradOptions(0.01) + .lr_decay(0) // learning rate decay + .weight_decay(0) + .initial_accumulator_value(0)); +``` diff --git a/docs/cpp/source/api/optim/gradient_descent.md b/docs/cpp/source/api/optim/gradient_descent.md new file mode 100644 index 0000000000000..c72b30171f717 --- /dev/null +++ b/docs/cpp/source/api/optim/gradient_descent.md @@ -0,0 +1,49 @@ +--- +myst: + html_meta: + description: SGD optimizer in PyTorch C++ — stochastic gradient descent with momentum and weight decay. + keywords: PyTorch, C++, SGD, gradient descent, momentum, weight decay, optimizer +--- + +# Gradient Descent Optimizers + +These optimizers use gradient descent with optional enhancements like momentum. +They are the foundation of neural network training and work well when you can +afford careful hyperparameter tuning. + +## SGD (Stochastic Gradient Descent) + +The classic optimization algorithm. SGD with momentum is often the best choice +for convolutional neural networks when properly tuned. While requiring more +careful learning rate selection than adaptive methods, it frequently achieves +the best final accuracy. + +**When to use:** + +- Training CNNs (ResNet, VGG, etc.) where you want maximum accuracy +- When you have time for hyperparameter tuning +- When combined with learning rate schedules (warmup, cosine annealing) + +**Key parameters:** + +- `lr`: Learning rate (typical: 0.01-0.1 for CNNs) +- `momentum`: Accelerates convergence (typical: 0.9) +- `weight_decay`: L2 regularization coefficient +- `nesterov`: Use Nesterov momentum (often improves convergence) + +```{doxygenclass} torch::optim::SGD +:members: +:undoc-members: +``` + +**Example:** + +```cpp +// Standard SGD with momentum - good for CNNs +auto optimizer = torch::optim::SGD( + model->parameters(), + torch::optim::SGDOptions(0.01) // learning rate + .momentum(0.9) // momentum factor + .weight_decay(1e-4) // L2 regularization + .nesterov(true)); // Nesterov momentum +``` diff --git a/docs/cpp/source/api/optim/index.md b/docs/cpp/source/api/optim/index.md new file mode 100644 index 0000000000000..e46c75be09f55 --- /dev/null +++ b/docs/cpp/source/api/optim/index.md @@ -0,0 +1,116 @@ +--- +myst: + html_meta: + description: PyTorch C++ optimizer API — SGD, Adam, and other optimizers for training neural networks. + keywords: PyTorch, C++, optimizer, optim, SGD, Adam, training +--- + +# Optimizers (torch::optim) + +The `torch::optim` namespace provides optimization algorithms for +training neural networks. These optimizers update model parameters based +on computed gradients to minimize the loss function. + +**When to use torch::optim:** + +- When training neural networks with gradient descent +- When you need different optimization strategies (SGD, Adam, etc.) +- When implementing learning rate schedules + +**Basic usage:** + +```cpp +#include + +// Create model and optimizer +auto model = std::make_shared(); +auto optimizer = torch::optim::Adam( + model->parameters(), + torch::optim::AdamOptions(1e-3)); + +// Training loop +for (auto& batch : *data_loader) { + optimizer.zero_grad(); // Clear gradients + auto loss = loss_fn(model->forward(batch.data), batch.target); + loss.backward(); // Compute gradients + optimizer.step(); // Update parameters +} +``` + +## Header Files + +- `torch/csrc/api/include/torch/optim.h` - Main optim header +- `torch/csrc/api/include/torch/optim/optimizer.h` - Optimizer base class +- `torch/csrc/api/include/torch/optim/sgd.h` - SGD optimizer +- `torch/csrc/api/include/torch/optim/adam.h` - Adam optimizer + +## Optimizer Base Class + +All optimizers inherit from the `Optimizer` base class, which provides common +functionality for parameter updates, gradient zeroing, and state management. + +```{doxygenclass} torch::optim::Optimizer +:members: +:undoc-members: +``` + +### OptimizerOptions + +```{doxygenclass} torch::optim::OptimizerOptions +:members: +:undoc-members: +``` + +### OptimizerParamGroup + +```{doxygenclass} torch::optim::OptimizerParamGroup +:members: +:undoc-members: +``` + +### OptimizerParamState + +```{doxygenclass} torch::optim::OptimizerParamState +:members: +:undoc-members: +``` + +## Choosing an Optimizer + +Selecting the right optimizer depends on your model architecture, dataset, and +training requirements: + +```{list-table} +:widths: 20 40 40 +:header-rows: 1 + +* - Optimizer + - Best For + - Trade-offs +* - **SGD + Momentum** + - CNNs, well-understood problems, when you can tune hyperparameters + - Requires careful learning rate tuning; often achieves best final accuracy +* - **Adam/AdamW** + - General-purpose, transformers, quick prototyping + - Works well out-of-the-box; AdamW preferred with weight decay +* - **RMSprop** + - RNNs, non-stationary objectives + - Good for recurrent architectures; handles varying gradient scales +* - **Adagrad** + - Sparse data (NLP, embeddings) + - Learning rate decreases over time; good for infrequent features +* - **LBFGS** + - Small models, fine-tuning, convex problems + - Memory-intensive; requires closure function +``` + +## Optimizer Categories + +```{toctree} +:maxdepth: 1 + +gradient_descent +adaptive +second_order +schedulers +``` diff --git a/docs/cpp/source/api/optim/schedulers.md b/docs/cpp/source/api/optim/schedulers.md new file mode 100644 index 0000000000000..814ff45e2fbaf --- /dev/null +++ b/docs/cpp/source/api/optim/schedulers.md @@ -0,0 +1,152 @@ +--- +myst: + html_meta: + description: Learning rate schedulers in PyTorch C++ — StepLR, ExponentialLR, and other LR scheduling policies. + keywords: PyTorch, C++, learning rate, scheduler, StepLR, ExponentialLR, LRScheduler +--- + +# Learning Rate Schedulers + +Learning rate schedulers adjust the learning rate during training, which often +improves convergence and final accuracy. Common strategies include: + +- **Step decay**: Reduce LR by a factor every N epochs +- **Exponential decay**: Multiply LR by gamma each epoch +- **Cosine annealing**: Smoothly decrease LR following a cosine curve +- **Warmup**: Gradually increase LR at the start of training + +## LRScheduler Base Class + +```{doxygenclass} torch::optim::LRScheduler +:members: +:undoc-members: +``` + +## StepLR + +Decays the learning rate by `gamma` every `step_size` epochs. This is the +simplest and most commonly used scheduler. + +```{doxygenclass} torch::optim::StepLR +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto optimizer = torch::optim::SGD( + model->parameters(), + torch::optim::SGDOptions(0.1)); + +// Reduce LR by 10x every 30 epochs +auto scheduler = torch::optim::StepLR( + optimizer, + /*step_size=*/30, + /*gamma=*/0.1); + +for (int epoch = 0; epoch < 90; ++epoch) { + train_one_epoch(model, optimizer, data_loader); + scheduler.step(); + // LR: 0.1 (epochs 0-29), 0.01 (30-59), 0.001 (60-89) +} +``` + +## ReduceLROnPlateau + +Reduces the learning rate when a metric has stopped improving. Useful when +you want the scheduler to respond to validation loss rather than follow a +fixed schedule. + +```{doxygenclass} torch::optim::ReduceLROnPlateauScheduler +:members: +:undoc-members: +``` + +## ExponentialLR + +Decays the learning rate by `gamma` every epoch. Provides smoother decay than +StepLR but may be slower to reduce the learning rate. + +**Example:** + +```cpp +auto optimizer = torch::optim::Adam( + model->parameters(), + torch::optim::AdamOptions(1e-3)); + +// Reduce LR by 5% each epoch +auto scheduler = torch::optim::ExponentialLR( + optimizer, + /*gamma=*/0.95); + +for (int epoch = 0; epoch < num_epochs; ++epoch) { + train_one_epoch(model, optimizer, data_loader); + scheduler.step(); +} +``` + +## Complete Training Example + +Here's a complete example showing optimizer usage with learning rate scheduling: + +```cpp +#include + +struct Net : torch::nn::Module { + Net() { + fc1 = register_module("fc1", torch::nn::Linear(784, 256)); + fc2 = register_module("fc2", torch::nn::Linear(256, 10)); + } + + torch::Tensor forward(torch::Tensor x) { + x = torch::relu(fc1->forward(x.view({-1, 784}))); + return fc2->forward(x); + } + + torch::nn::Linear fc1{nullptr}, fc2{nullptr}; +}; + +int main() { + // Create model + auto model = std::make_shared(); + + // Create optimizer with weight decay + auto optimizer = torch::optim::AdamW( + model->parameters(), + torch::optim::AdamWOptions(1e-3) + .weight_decay(0.01)); + + // Learning rate scheduler + auto scheduler = torch::optim::StepLR(optimizer, 10, 0.5); + + // Loss function + auto loss_fn = torch::nn::CrossEntropyLoss(); + + // Training loop + for (int epoch = 0; epoch < 30; ++epoch) { + model->train(); + double epoch_loss = 0.0; + + for (auto& batch : *train_loader) { + optimizer.zero_grad(); + + auto output = model->forward(batch.data); + auto loss = loss_fn(output, batch.target); + + loss.backward(); + optimizer.step(); + + epoch_loss += loss.item(); + } + + scheduler.step(); + std::cout << "Epoch " << epoch + << " Loss: " << epoch_loss + << " LR: " << scheduler.get_last_lr()[0] + << std::endl; + } + + return 0; +} +``` diff --git a/docs/cpp/source/api/optim/second_order.md b/docs/cpp/source/api/optim/second_order.md new file mode 100644 index 0000000000000..0f93621aecf41 --- /dev/null +++ b/docs/cpp/source/api/optim/second_order.md @@ -0,0 +1,60 @@ +--- +myst: + html_meta: + description: Second-order optimizers in PyTorch C++ — LBFGS optimizer for full-batch optimization. + keywords: PyTorch, C++, LBFGS, second-order, optimizer, full-batch +--- + +# Second-Order Optimizers + +Second-order methods use curvature information (Hessian or its approximations) +to make better optimization steps. They can converge faster but are more +computationally expensive and memory-intensive. + +## LBFGS (Limited-memory Broyden-Fletcher-Goldfarb-Shanno) + +LBFGS is a quasi-Newton method that approximates the inverse Hessian using +gradient history. It can converge much faster than first-order methods for +smooth, convex-like loss surfaces. + +**When to use:** + +- Small models where memory isn't a concern +- Fine-tuning pre-trained models +- Convex or near-convex optimization problems +- Full-batch training (not mini-batch) + +**Key parameters:** + +- `lr`: Learning rate (often 1.0 for LBFGS) +- `max_iter`: Maximum iterations per step +- `history_size`: Number of past gradients to store + +**Important:** LBFGS requires a closure function that recomputes the loss. + +```{doxygenclass} torch::optim::LBFGS +:members: +:undoc-members: +``` + +**Example:** + +```cpp +auto optimizer = torch::optim::LBFGS( + model->parameters(), + torch::optim::LBFGSOptions(1.0) + .max_iter(20) + .history_size(10)); + +// LBFGS requires a closure that recomputes the model +for (int epoch = 0; epoch < num_epochs; ++epoch) { + auto closure = [&]() { + optimizer.zero_grad(); + auto output = model->forward(data); + auto loss = loss_fn(output, target); + loss.backward(); + return loss; + }; + optimizer.step(closure); +} +``` diff --git a/docs/cpp/source/api/serialize/archives.md b/docs/cpp/source/api/serialize/archives.md new file mode 100644 index 0000000000000..2d50e3db08c9f --- /dev/null +++ b/docs/cpp/source/api/serialize/archives.md @@ -0,0 +1,72 @@ +--- +myst: + html_meta: + description: Serialization archives in PyTorch C++ — InputArchive and OutputArchive for model serialization. + keywords: PyTorch, C++, InputArchive, OutputArchive, serialization, archive +--- + +# Archives + +Archives provide a lower-level interface for serialization, allowing you to +save multiple values to a single file with named keys. + +## OutputArchive + +```{doxygenclass} torch::serialize::OutputArchive +:members: +:undoc-members: +``` + +**Example:** + +```cpp +torch::serialize::OutputArchive archive; +archive.write("tensor1", tensor1); +archive.write("tensor2", tensor2); +archive.save_to("multi_tensor.pt"); +``` + +## InputArchive + +```{doxygenclass} torch::serialize::InputArchive +:members: +:undoc-members: +``` + +**Example:** + +```cpp +torch::serialize::InputArchive archive; +archive.load_from("multi_tensor.pt"); + +torch::Tensor tensor1, tensor2; +archive.read("tensor1", tensor1); +archive.read("tensor2", tensor2); +``` + +## Saving Multiple Values + +Archives are useful when you need to save multiple related values together: + +```cpp +// Save multiple tensors and metadata +torch::serialize::OutputArchive out_archive; +out_archive.write("weights", model_weights); +out_archive.write("biases", model_biases); +out_archive.write("epoch", torch::tensor(current_epoch)); +out_archive.write("loss", torch::tensor(best_loss)); +out_archive.save_to("training_state.pt"); + +// Load them back +torch::serialize::InputArchive in_archive; +in_archive.load_from("training_state.pt"); + +torch::Tensor weights, biases, epoch_tensor, loss_tensor; +in_archive.read("weights", weights); +in_archive.read("biases", biases); +in_archive.read("epoch", epoch_tensor); +in_archive.read("loss", loss_tensor); + +int epoch = epoch_tensor.item(); +float loss = loss_tensor.item(); +``` diff --git a/docs/cpp/source/api/serialize/checkpoints.md b/docs/cpp/source/api/serialize/checkpoints.md new file mode 100644 index 0000000000000..27cae975b47c5 --- /dev/null +++ b/docs/cpp/source/api/serialize/checkpoints.md @@ -0,0 +1,113 @@ +--- +myst: + html_meta: + description: Checkpointing in PyTorch C++ — saving and resuming training state. + keywords: PyTorch, C++, checkpoint, save, resume, training state +--- + +# Checkpoints + +Checkpoints save the complete training state so you can resume training +after interruption. A checkpoint typically includes: + +- Model parameters +- Optimizer state (momentum buffers, learning rates) +- Current epoch number +- Best validation loss/accuracy + +## Creating Checkpoints + +```cpp +void save_checkpoint( + std::shared_ptr model, + torch::optim::Adam& optimizer, + int epoch, + const std::string& path) { + torch::serialize::OutputArchive archive; + model->save(archive); + archive.write("epoch", torch::tensor(epoch)); + optimizer.save(archive); + archive.save_to(path); +} +``` + +## Loading Checkpoints + +```cpp +int load_checkpoint( + std::shared_ptr model, + torch::optim::Adam& optimizer, + const std::string& path) { + torch::serialize::InputArchive archive; + archive.load_from(path); + model->load(archive); + torch::Tensor epoch_tensor; + archive.read("epoch", epoch_tensor); + optimizer.load(archive); + return epoch_tensor.item(); +} +``` + +## Complete Checkpoint Example + +```cpp +#include +#include +#include + +struct Net : torch::nn::Module { + Net() { + fc1 = register_module("fc1", torch::nn::Linear(784, 256)); + fc2 = register_module("fc2", torch::nn::Linear(256, 10)); + } + + torch::Tensor forward(torch::Tensor x) { + x = torch::relu(fc1->forward(x.view({-1, 784}))); + return fc2->forward(x); + } + + torch::nn::Linear fc1{nullptr}, fc2{nullptr}; +}; + +int main() { + auto model = std::make_shared(); + auto optimizer = torch::optim::Adam(model->parameters(), 1e-3); + + int start_epoch = 0; + const std::string checkpoint_path = "checkpoint.pt"; + + // Resume from checkpoint if it exists + if (std::filesystem::exists(checkpoint_path)) { + std::cout << "Loading checkpoint..." << std::endl; + start_epoch = load_checkpoint(model, optimizer, checkpoint_path); + std::cout << "Resuming from epoch " << start_epoch << std::endl; + } + + // Training loop + for (int epoch = start_epoch; epoch < 100; ++epoch) { + // ... training code ... + + // Save checkpoint every 10 epochs + if ((epoch + 1) % 10 == 0) { + save_checkpoint(model, optimizer, epoch + 1, checkpoint_path); + std::cout << "Saved checkpoint at epoch " << epoch + 1 << std::endl; + } + } + + return 0; +} +``` + +## Best Practices + +1. **Save periodically**: Save checkpoints at regular intervals (e.g., every epoch + or every N batches) to minimize lost work. + +2. **Keep multiple checkpoints**: Maintain the last few checkpoints in case the + most recent one is corrupted or represents a poor model state. + +3. **Include all state**: Save everything needed to resume training, including + learning rate scheduler state if using one. + +4. **Verify checkpoints**: Occasionally verify that checkpoints can be loaded + correctly. diff --git a/docs/cpp/source/api/serialize/index.md b/docs/cpp/source/api/serialize/index.md new file mode 100644 index 0000000000000..002431d9ef632 --- /dev/null +++ b/docs/cpp/source/api/serialize/index.md @@ -0,0 +1,58 @@ +--- +myst: + html_meta: + description: PyTorch C++ serialization API — saving and loading models and tensors. + keywords: PyTorch, C++, serialization, save, load, checkpoint, model +--- + +# Serialization (torch::serialize) + +The `torch::serialize` namespace provides utilities for saving and loading +model weights, tensors, and optimizer state. This enables checkpointing during +training and deployment of trained models. + +**When to use torch::serialize:** + +- When saving trained models to disk +- When implementing training checkpoints +- When loading pre-trained weights +- When transferring models between C++ and Python + +**Basic usage:** + +```cpp +#include + +// Save a model +auto model = std::make_shared(); +// ... train the model ... +torch::save(model, "model.pt"); + +// Load a model +auto loaded_model = std::make_shared(); +torch::load(loaded_model, "model.pt"); + +// Save and load tensors +torch::Tensor tensor = torch::randn({2, 3}); +torch::save(tensor, "tensor.pt"); + +torch::Tensor loaded_tensor; +torch::load(loaded_tensor, "tensor.pt"); +``` + +## Header Files + +- `torch/csrc/api/include/torch/serialize.h` - Main serialization header +- `torch/csrc/api/include/torch/serialize/archive.h` - Archive classes +- `torch/csrc/api/include/torch/serialize/input-archive.h` - Input archive +- `torch/csrc/api/include/torch/serialize/output-archive.h` - Output archive + +## Serialization Categories + +```{toctree} +:maxdepth: 1 + +save_load +archives +checkpoints +``` diff --git a/docs/cpp/source/api/serialize/save_load.md b/docs/cpp/source/api/serialize/save_load.md new file mode 100644 index 0000000000000..e08cd21f905cd --- /dev/null +++ b/docs/cpp/source/api/serialize/save_load.md @@ -0,0 +1,84 @@ +--- +myst: + html_meta: + description: Save and load functions in PyTorch C++ — torch::save and torch::load for tensors and modules. + keywords: PyTorch, C++, save, load, torch::save, torch::load, tensor, module +--- + +# Saving and Loading + +The primary interface for serialization uses the `torch::save` and +`torch::load` functions, which can save and load tensors, modules, +and optimizers. + +## Save Functions + +```{doxygenfunction} torch::save(const Value &value, SaveToArgs&&... args) +``` + +```{doxygenfunction} torch::save(const std::vector &tensor_vec, SaveToArgs&&... args) +``` + +## Load Functions + +```{doxygenfunction} torch::load(Value &value, LoadFromArgs&&... args) +``` + +```{doxygenfunction} torch::load(std::vector &tensor_vec, LoadFromArgs&&... args) +``` + +## Saving and Loading Tensors + +```cpp +// Save a tensor +torch::Tensor tensor = torch::randn({2, 3}); +torch::save(tensor, "tensor.pt"); + +// Load a tensor +torch::Tensor loaded; +torch::load(loaded, "tensor.pt"); +``` + +## Saving and Loading Modules + +```cpp +// Define a model +struct Net : torch::nn::Module { + Net() { + fc1 = register_module("fc1", torch::nn::Linear(784, 64)); + fc2 = register_module("fc2", torch::nn::Linear(64, 10)); + } + + torch::Tensor forward(torch::Tensor x) { + x = torch::relu(fc1->forward(x)); + return fc2->forward(x); + } + + torch::nn::Linear fc1{nullptr}, fc2{nullptr}; +}; + +// Save model +auto model = std::make_shared(); +torch::save(model, "model.pt"); + +// Load model +auto loaded_model = std::make_shared(); +torch::load(loaded_model, "model.pt"); +``` + +## Saving Optimizer State + +```cpp +auto model = std::make_shared(); +auto optimizer = torch::optim::Adam(model->parameters(), 0.001); + +// Train... + +// Save both model and optimizer +torch::save(model, "model.pt"); +torch::save(optimizer, "optimizer.pt"); + +// Load both +torch::load(model, "model.pt"); +torch::load(optimizer, "optimizer.pt"); +``` diff --git a/docs/cpp/source/api/stable/index.md b/docs/cpp/source/api/stable/index.md new file mode 100644 index 0000000000000..adb3d477d18b6 --- /dev/null +++ b/docs/cpp/source/api/stable/index.md @@ -0,0 +1,72 @@ +--- +myst: + html_meta: + description: PyTorch Stable ABI C++ API — binary-compatible operator registration across PyTorch versions. + keywords: PyTorch, C++, stable ABI, binary compatibility, operator registration +--- + +# Torch Stable API + +The PyTorch Stable C++ API provides a binary-compatible interface for calling +tensor operations and utilities that is guaranteed to remain stable across +PyTorch versions. This enables ahead-of-time compiled extensions that don't +need recompilation when PyTorch is updated. + +**When to use the Stable API:** + +- When building extensions that must work across multiple PyTorch versions +- When distributing pre-compiled binaries +- When binary compatibility is more important than access to the latest features +- When writing custom operators for production deployment + +**Basic usage:** + +```cpp +#include +#include + +// Create a tensor using stable API +auto tensor = torch::stable::empty( + {3, 4}, + torch::headeronly::ScalarType::Float, + torch::headeronly::Layout::Strided, + torch::stable::Device(torch::headeronly::DeviceType::CPU), + false, + torch::headeronly::MemoryFormat::Contiguous); + +// Register operators with stable ABI +STABLE_TORCH_LIBRARY(myops, m) { + m.def("my_op(Tensor input) -> Tensor"); +} + +STABLE_TORCH_LIBRARY_IMPL(myops, CPU, m) { + m.impl("my_op", TORCH_BOX(&my_cpu_kernel)); +} +``` + +For more information on the stable ABI, see the +[Stable ABI notes](https://docs.pytorch.org/docs/stable/notes/libtorch_stable_abi.html). + +## Header Files + +- `torch/csrc/stable/library.h` - Stable library registration +- `torch/csrc/stable/ops.h` - Stable operator definitions +- `torch/csrc/stable/tensor_struct.h` - Stable tensor structures +- `torch/csrc/stable/device_struct.h` - Stable device structures +- `torch/csrc/stable/accelerator.h` - Accelerator support +- `torch/csrc/stable/macros.h` - Stable API macros + +## Stable API Categories + +```{toctree} +:maxdepth: 1 + +registration +operators +utilities +``` + +## See Also + +- {doc}`../library/index` - For standard (non-stable) operator registration +- [Stable ABI documentation](https://pytorch.org/docs/stable/cpp_extension.html) diff --git a/docs/cpp/source/api/stable/operators.md b/docs/cpp/source/api/stable/operators.md new file mode 100644 index 0000000000000..8c3bfac176ea6 --- /dev/null +++ b/docs/cpp/source/api/stable/operators.md @@ -0,0 +1,146 @@ +--- +myst: + html_meta: + description: Stable ABI operator API in PyTorch C++ — StableLibrary and boxed kernel registration. + keywords: PyTorch, C++, stable ABI, StableLibrary, operator, boxed kernel +--- + +# Stable Operators + +The stable API provides tensor operations that maintain binary compatibility +across PyTorch versions. + +## Tensor Class + +```{doxygenclass} torch::stable::Tensor +:members: +:undoc-members: +``` + +**Example:** + +```cpp +torch::stable::Tensor tensor = torch::stable::empty({3, 4}, ...); +float* data = tensor.data_ptr(); +auto shape = tensor.sizes(); +``` + +## Device Class + +```{doxygenclass} torch::stable::Device +:members: +:undoc-members: +``` + +**Example:** + +```cpp +torch::stable::Device cpu_device(torch::headeronly::DeviceType::CPU); +torch::stable::Device cuda_device(torch::headeronly::DeviceType::CUDA, 0); +``` + +## Tensor Creation + +```{doxygenfunction} torch::stable::empty +``` + +```{doxygenfunction} torch::stable::empty_like +``` + +```{doxygenfunction} torch::stable::new_empty(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef size, std::optional dtype, std::optional layout, std::optional device, std::optional pin_memory) +``` + +```{doxygenfunction} torch::stable::new_zeros(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef size, std::optional dtype, std::optional layout, std::optional device, std::optional pin_memory) +``` + +```{doxygenfunction} torch::stable::full +``` + +```{doxygenfunction} torch::stable::from_blob(void *data, torch::headeronly::IntHeaderOnlyArrayRef sizes, torch::headeronly::IntHeaderOnlyArrayRef strides, torch::stable::Device device, torch::headeronly::ScalarType dtype, int64_t storage_offset, torch::headeronly::Layout layout) +``` + +**Example:** + +```cpp +auto tensor = torch::stable::empty( + {3, 4}, + torch::headeronly::ScalarType::Float, + torch::headeronly::Layout::Strided, + torch::stable::Device(torch::headeronly::DeviceType::CUDA, 0), + false, + torch::headeronly::MemoryFormat::Contiguous); +``` + +## Tensor Manipulation + +```{doxygenfunction} torch::stable::clone +``` + +```{doxygenfunction} torch::stable::contiguous +``` + +```{doxygenfunction} torch::stable::reshape +``` + +```{doxygenfunction} torch::stable::view +``` + +```{doxygenfunction} torch::stable::flatten +``` + +```{doxygenfunction} torch::stable::squeeze +``` + +```{doxygenfunction} torch::stable::unsqueeze +``` + +```{doxygenfunction} torch::stable::transpose +``` + +```{doxygenfunction} torch::stable::select +``` + +```{doxygenfunction} torch::stable::narrow +``` + +```{doxygenfunction} torch::stable::pad +``` + +## Device and Type Conversion + +```{doxygenfunction} torch::stable::to(const torch::stable::Tensor &self, std::optional dtype, std::optional layout, std::optional device, std::optional pin_memory, bool non_blocking, bool copy, std::optional memory_format) +``` + +```{doxygenfunction} torch::stable::to(const torch::stable::Tensor &self, torch::stable::Device device, bool non_blocking, bool copy) +``` + +## In-place Operations + +```{doxygenfunction} torch::stable::fill_ +``` + +```{doxygenfunction} torch::stable::zero_ +``` + +```{doxygenfunction} torch::stable::copy_ +``` + +## Mathematical Operations + +```{doxygenfunction} torch::stable::matmul +``` + +```{doxygenfunction} torch::stable::amax(const torch::stable::Tensor &self, int64_t dim, bool keepdim) +``` + +```{doxygenfunction} torch::stable::amax(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef dims, bool keepdim) +``` + +```{doxygenfunction} torch::stable::sum +``` + +```{doxygenfunction} torch::stable::sum_out +``` + +```{doxygenfunction} torch::stable::subtract +``` diff --git a/docs/cpp/source/api/stable/registration.md b/docs/cpp/source/api/stable/registration.md new file mode 100644 index 0000000000000..ad970a66d72ae --- /dev/null +++ b/docs/cpp/source/api/stable/registration.md @@ -0,0 +1,116 @@ +--- +myst: + html_meta: + description: Stable ABI registration macros in PyTorch C++ — STABLE_TORCH_LIBRARY, STABLE_TORCH_LIBRARY_IMPL, and TORCH_BOX. + keywords: PyTorch, C++, stable ABI, STABLE_TORCH_LIBRARY, TORCH_BOX, macro, registration +--- + +# Library Registration Macros + +These macros provide stable ABI equivalents of the standard PyTorch operator +registration macros (`TORCH_LIBRARY`, `TORCH_LIBRARY_IMPL`, etc.). +Use these when building custom operators that need to maintain binary +compatibility across PyTorch versions. + +## STABLE_TORCH_LIBRARY + +```{c:macro} STABLE_TORCH_LIBRARY(ns, m) + +Defines a library of operators in a namespace using the stable ABI. + +This is the stable ABI equivalent of {c:macro}`TORCH_LIBRARY`. +Use this macro to define operator schemas that will maintain +binary compatibility across PyTorch versions. Only one `STABLE_TORCH_LIBRARY` +block can exist per namespace; use `STABLE_TORCH_LIBRARY_FRAGMENT` for +additional definitions in the same namespace from different translation units. + +:param ns: The namespace in which to define operators (e.g., `mylib`). +:param m: The name of the StableLibrary variable available in the block. + +**Example:** + +```cpp +STABLE_TORCH_LIBRARY(mylib, m) { + m.def("my_op(Tensor input, int size) -> Tensor"); + m.def("another_op(Tensor a, Tensor b) -> Tensor"); +} +``` + +Minimum compatible version: PyTorch 2.9. +``` +## STABLE_TORCH_LIBRARY_IMPL + +```{c:macro} STABLE_TORCH_LIBRARY_IMPL(ns, k, m) + +Registers operator implementations for a specific dispatch key using the stable ABI. + +This is the stable ABI equivalent of `TORCH_LIBRARY_IMPL`. Use this macro +to provide implementations of operators for a specific dispatch key (e.g., +CPU, CUDA) while maintaining binary compatibility across PyTorch versions. +``` + +```{note} + +All kernel functions registered with this macro must be boxed using +the `TORCH_BOX` macro. +``` + +:param ns: The namespace in which the operators are defined. +:param k: The dispatch key (e.g., `CPU`, `CUDA`). +:param m: The name of the StableLibrary variable available in the block. + +**Example:** + +```cpp +STABLE_TORCH_LIBRARY_IMPL(mylib, CPU, m) { + m.impl("my_op", TORCH_BOX(&my_cpu_kernel)); +} + +STABLE_TORCH_LIBRARY_IMPL(mylib, CUDA, m) { + m.impl("my_op", TORCH_BOX(&my_cuda_kernel)); +} +``` + +Minimum compatible version: PyTorch 2.9. + +## STABLE_TORCH_LIBRARY_FRAGMENT + +```{c:macro} STABLE_TORCH_LIBRARY_FRAGMENT(ns, m) + +Extends operator definitions in an existing namespace using the stable ABI. + +This is the stable ABI equivalent of `TORCH_LIBRARY_FRAGMENT`. Use this macro +to add additional operator definitions to a namespace that was already +created with `STABLE_TORCH_LIBRARY`. + +:param ns: The namespace to extend. +:param m: The name of the StableLibrary variable available in the block. + +Minimum compatible version: PyTorch 2.9. +``` + +## TORCH_BOX + +```{c:macro} TORCH_BOX(func) + +Wraps a function to conform to the stable boxed kernel calling convention. + +This macro takes an unboxed kernel function pointer and generates a boxed wrapper +that can be registered with the stable library API. + +:param func: The unboxed kernel function to wrap. + +**Example:** + +```cpp +Tensor my_kernel(const Tensor& input, int64_t size) { + return input.reshape({size}); +} + +STABLE_TORCH_LIBRARY_IMPL(my_namespace, CPU, m) { + m.impl("my_op", TORCH_BOX(&my_kernel)); +} +``` + +Minimum compatible version: PyTorch 2.9. +``` \ No newline at end of file diff --git a/docs/cpp/source/api/stable/utilities.md b/docs/cpp/source/api/stable/utilities.md new file mode 100644 index 0000000000000..c9d30e96ddef9 --- /dev/null +++ b/docs/cpp/source/api/stable/utilities.md @@ -0,0 +1,267 @@ +--- +myst: + html_meta: + description: Stable ABI utilities in PyTorch C++ — version checking and compatibility helpers. + keywords: PyTorch, C++, stable ABI, utilities, version, compatibility +--- + +# Utilities + +The stable API provides various utility functions and types for working with +tensors and CUDA operations. + +## DeviceGuard Class + +```{doxygenclass} torch::stable::accelerator::DeviceGuard +:members: +:undoc-members: +``` + +```{doxygenfunction} torch::stable::accelerator::getCurrentDeviceIndex +``` + +**Example:** + +```cpp +{ + torch::stable::accelerator::DeviceGuard guard(1); + // Operations here run on device 1 +} +// Previous device is restored +``` + +## Stream + +```{doxygenclass} torch::stable::accelerator::Stream +:members: +:undoc-members: +``` + +## Stream Utilities + +For CUDA stream access, we currently recommend the ABI stable C shim API. This +will be improved in a future release with a more ergonomic wrapper. + +### Getting the Current CUDA Stream + +To obtain the current `cudaStream_t` for use in CUDA kernels: + +```cpp +#include +#include + +// For now, we rely on the ABI stable C shim API to get the current CUDA stream. +void* stream_ptr = nullptr; +TORCH_ERROR_CODE_CHECK( + aoti_torch_get_current_cuda_stream(tensor.get_device_index(), &stream_ptr)); +cudaStream_t stream = static_cast(stream_ptr); + +// Now you can use 'stream' in your CUDA kernel launches +my_kernel<<>>(args...); +``` + +```{note} + +The `TORCH_ERROR_CODE_CHECK` macro is required when using C shim APIs +to properly check error codes and throw appropriate exceptions. +``` + +## CUDA Error Checking Macros + +These macros provide stable ABI equivalents for CUDA error checking. +They wrap CUDA API calls and kernel launches, providing detailed error +messages using PyTorch's error formatting. + +### STD_CUDA_CHECK + +```{c:macro} STD_CUDA_CHECK(EXPR) + +Checks the result of a CUDA API call and throws an exception on error. +Users of this macro are expected to include `cuda_runtime.h`. + +**Example:** + +```cpp +STD_CUDA_CHECK(cudaMalloc(&ptr, size)); +STD_CUDA_CHECK(cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost)); +``` + +Minimum compatible version: PyTorch 2.10. +``` +### STD_CUDA_KERNEL_LAUNCH_CHECK + +```{c:macro} STD_CUDA_KERNEL_LAUNCH_CHECK() + +Checks for errors from the most recent CUDA kernel launch. Equivalent to +`STD_CUDA_CHECK(cudaGetLastError())`. + +**Example:** + +```cpp +my_kernel<<>>(args...); +STD_CUDA_KERNEL_LAUNCH_CHECK(); +``` + +Minimum compatible version: PyTorch 2.10. +``` +## Header-Only Utilities + +The `torch::headeronly` namespace provides header-only versions of common +PyTorch types and utilities. These can be used without linking against libtorch +at all! This portability makes them ideal for maintaining binary compatibility +across PyTorch versions. + +### Error Checking + +`STD_TORCH_CHECK` is a header-only macro for runtime assertions: + +```cpp +#include + +STD_TORCH_CHECK(condition, "Error message with ", variable, " interpolation"); +``` + +Wherever you used `TORCH_CHECK` before, you can replace usage with `STD_TORCH_CHECK` +to remove the need to link against libtorch. The only difference is that when the +condition check fails, `TORCH_CHECK` throws a fancier `c10::Error` while +`STD_TORCH_CHECK` throws a `std::runtime_error`. + +### Core Types + +The following `c10::` types are available as header-only versions under +`torch::headeronly::`: + +- `torch::headeronly::ScalarType` - Tensor data types (Float, Double, Int, etc.) +- `torch::headeronly::DeviceType` - Device types (CPU, CUDA, etc.) +- `torch::headeronly::MemoryFormat` - Memory layout formats (Contiguous, ChannelsLast, etc.) +- `torch::headeronly::Layout` - Tensor layouts (Strided, Sparse, etc.) + +```cpp +#include +#include +#include +#include + +auto dtype = torch::headeronly::ScalarType::Float; +auto device_type = torch::headeronly::DeviceType::CUDA; +auto memory_format = torch::headeronly::MemoryFormat::Contiguous; +auto layout = torch::headeronly::Layout::Strided; +``` + +### TensorAccessor + +`TensorAccessor` provides efficient, bounds-checked access to tensor data. +You can construct one from a stable tensor's data pointer, sizes, and strides: + +```cpp +#include + +// Create a TensorAccessor for a 2D float tensor +auto sizes = tensor.sizes(); +auto strides = tensor.strides(); +torch::headeronly::TensorAccessor accessor( + static_cast(tensor.mutable_data_ptr()), + sizes.data(), + strides.data()); + +// Access elements +float value = accessor[i][j]; +``` + +### Dispatch Macros + +Header-only dispatch macros (THO = Torch Header Only) are available for +dtype dispatching: + +```cpp +#include + +THO_DISPATCH_V2( + tensor.scalar_type(), // will be resolved as scalar_t + "my_kernel", + AT_WRAP(([&]() { + // code to specialize with scalar_t + // scalar_t is the resolved C++ type (e.g. float, double) + auto* data = static_cast(tensor.mutable_data_ptr()); + Scalar s(*data); + })), + AT_EXPAND(AT_ALL_TYPES), + AT_EXPAND(AT_COMPLEX_TYPES), + torch::headeronly::ScalarType::Half, + // as many type arguments as needed +); +``` + +`THO_DISPATCH_V2` works the same way as `AT_DISPATCH_V2` (see +`ATen/Dispatch_v2.h`) but does not require linking against libtorch. +As a result, whereas `AT_DISPATCH_V2` would have thrown `c10::NotImplementedError` +for unimplemented paths, `THO_DISPATCH_V2` will throw `std::runtime_error`. + +For ease of use, we've also migrated the below AT_* macros representing +collections of types to be header-only and thus have no dependency on libtorch: + +- `AT_FLOATING_TYPES` +- `AT_INTEGRAL_TYPES` +- `AT_INTEGRAL_TYPES_V2` +- `AT_ALL_TYPES` +- `AT_COMPLEX_TYPES` +- `AT_ALL_TYPES_AND_COMPLEX` +- `AT_FLOAT8_TYPES` +- `AT_BAREBONES_UNSIGNED_TYPES` +- `AT_QINT_TYPES` + +If your extension uses our older AT_DISPATCH version 1 infrastructure, +you can also migrate to a header-only libtorch-free world without upgrading +everything to version 2. + +`THO_DISPATCH_SWITCH` and `THO_DISPATCH_CASE` are the header-only +equivalents of `AT_DISPATCH_SWITCH` and `AT_DISPATCH_CASE`. Similarly, +the only user-visible difference is the exception type on an unhandled dtype, +where the `AT_` version throws a `c10::NotImplementedError` and the `THO_` +version throws a `std::runtime_error`. + +The migration is pretty mechanical: + +- `AT_DISPATCH_SWITCH` → `THO_DISPATCH_SWITCH` +- `AT_DISPATCH_CASE` → `THO_DISPATCH_CASE` +- `AT_PRIVATE_CASE_TYPE_USING_HINT` → `THO_PRIVATE_CASE_TYPE_USING_HINT` +- `at::ScalarType::X` → `torch::headeronly::ScalarType::X` + +```cpp +// ---- Before (requires linking against libtorch) ---- +#include + +#define MY_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define MY_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, \ + MY_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) +``` + +```cpp +// ---- After (header-only, no libtorch dependency) ---- +#include + +#define MY_DISPATCH_CASE_FLOATING_TYPES(...) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \ + THO_DISPATCH_CASE(torch::headeronly::ScalarType::BFloat16, __VA_ARGS__) + +#define MY_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH(TYPE, NAME, \ + MY_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) +``` + +For the complete list of header-only APIs, see `torch/header_only_apis.txt` +in the PyTorch source tree. + +## Parallelization Utilities + +```{doxygenfunction} torch::stable::parallel_for +``` + +```{doxygenfunction} torch::stable::get_num_threads +``` diff --git a/docs/cpp/source/api/xpu/index.md b/docs/cpp/source/api/xpu/index.md new file mode 100644 index 0000000000000..c88eff6a69165 --- /dev/null +++ b/docs/cpp/source/api/xpu/index.md @@ -0,0 +1,48 @@ +--- +myst: + html_meta: + description: PyTorch XPU C++ API — Intel GPU support with device management, streams, and guards. + keywords: PyTorch, C++, XPU, Intel GPU, device, stream +--- + +# XPU Support + +PyTorch provides XPU support for Intel GPU-accelerated tensor operations. +The XPU API allows you to manage Intel GPU devices, streams for asynchronous +execution, and synchronization. + +**When to use XPU APIs:** + +- When running on Intel GPUs (Data Center GPU Max, Arc, etc.) +- When implementing custom XPU kernels or operations +- When managing asynchronous execution with XPU streams +- When writing device-portable code alongside CUDA + +**Basic usage:** + +```cpp +#include + +// Check if XPU is available +if (torch::xpu::is_available()) { + // Create tensor on XPU + auto tensor = torch::randn({2, 3}, torch::device(torch::kXPU)); + + // Move model to XPU + model->to(torch::kXPU); +} +``` + +## Header Files + +- `torch/xpu.h` - High-level XPU utilities (device count, availability, seeding) +- `c10/xpu/XPUStream.h` - XPU stream management + +## XPU Categories + +```{toctree} +:maxdepth: 1 + +streams +utilities +``` diff --git a/docs/cpp/source/api/xpu/streams.md b/docs/cpp/source/api/xpu/streams.md new file mode 100644 index 0000000000000..a9b6911471d4f --- /dev/null +++ b/docs/cpp/source/api/xpu/streams.md @@ -0,0 +1,50 @@ +--- +myst: + html_meta: + description: XPU streams in PyTorch C++ — XPUStream for asynchronous Intel GPU execution. + keywords: PyTorch, C++, XPU, XPUStream, Intel GPU, stream, asynchronous +--- + +# XPU Streams + +XPU streams provide a mechanism for asynchronous execution of operations +on Intel GPUs. Like CUDA streams, operations queued to the same stream execute +in order, while operations on different streams can execute concurrently. + +## XPUStream + +```{doxygenclass} c10::xpu::XPUStream +:members: +:undoc-members: +``` + +**Example:** + +```cpp +#include + +// Get the current XPU stream +auto stream = c10::xpu::getCurrentXPUStream(); + +// Create a new stream from the pool +auto new_stream = c10::xpu::getStreamFromPool(); + +// Synchronize +stream.synchronize(); +``` + +## Acquiring XPU Streams + +```{doxygenfunction} c10::xpu::getCurrentXPUStream +``` + +```{doxygenfunction} c10::xpu::setCurrentXPUStream +``` + +```{doxygenfunction} c10::xpu::getStreamFromPool(const bool isHighPriority, DeviceIndex device) +``` + +## Stream Synchronization + +```{doxygenfunction} c10::xpu::syncStreamsOnDevice +``` diff --git a/docs/cpp/source/api/xpu/utilities.md b/docs/cpp/source/api/xpu/utilities.md new file mode 100644 index 0000000000000..9857d5e291a2b --- /dev/null +++ b/docs/cpp/source/api/xpu/utilities.md @@ -0,0 +1,53 @@ +--- +myst: + html_meta: + description: XPU utility functions in PyTorch C++ — device count, properties, and stream management for Intel GPUs. + keywords: PyTorch, C++, XPU, Intel GPU, device_count, utilities +--- + +# XPU Utility Functions + +High-level utility functions for querying and managing XPU devices. + +## Device Management + +```{doxygenfunction} torch::xpu::device_count +``` + +```{doxygenfunction} torch::xpu::is_available +``` + +```{doxygenfunction} torch::xpu::synchronize +``` + +**Example:** + +```cpp +#include + +if (torch::xpu::is_available()) { + size_t num_devices = torch::xpu::device_count(); + std::cout << "Found " << num_devices << " XPU device(s)" << std::endl; + + // Synchronize all streams on device 0 + torch::xpu::synchronize(0); +} +``` + +## Random Number Generation + +```{doxygenfunction} torch::xpu::manual_seed +``` + +```{doxygenfunction} torch::xpu::manual_seed_all +``` + +**Example:** + +```cpp +// Set seed for reproducibility on current XPU device +torch::xpu::manual_seed(42); + +// Set seed for all XPU devices +torch::xpu::manual_seed_all(42); +``` diff --git a/docs/cpp/source/conf.py b/docs/cpp/source/conf.py index 10d854c21db4f..ab3c722bb63ed 100644 --- a/docs/cpp/source/conf.py +++ b/docs/cpp/source/conf.py @@ -20,7 +20,6 @@ # See https://github.com/pytorch/pytorch/issues/79992. import os -import textwrap # sys.path.insert(0, os.path.abspath('.')) import pytorch_sphinx_theme2 @@ -31,14 +30,55 @@ # If your documentation needs a minimal Sphinx version, state it here. # needs_sphinx = "3.1.2" -run_doxygen = os.environ.get("RUN_DOXYGEN", "false") == "true" # Add any Sphinx extension module names here, as strings. They can be # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = [ "sphinx.ext.intersphinx", -] + (["breathe", "exhale"] if run_doxygen else []) + "breathe", + "myst_parser", +] + +# -- Breathe Configuration ------------------------------------------------ +# Breathe connects Sphinx to Doxygen XML output +# Use doxygenclass, doxygenfunction, etc. directives in RST files +# to pull documentation from C++ source code + +breathe_projects = {"PyTorch": "../build/xml"} +breathe_default_project = "PyTorch" + +# Default members to show when using doxygenclass/doxygenstruct directives +breathe_default_members = ("members", "undoc-members") + +# Map file extensions to language domains for proper syntax highlighting +breathe_domain_by_extension = { + "h": "cpp", + "hpp": "cpp", + "cpp": "cpp", + "c": "c", +} + +# Implementation detail filters - skip internal/private content +breathe_implementation_filename_extensions = [".c", ".cc", ".cpp"] + +# Show the file where items are defined +breathe_show_define_initializer = True +breathe_show_enumvalue_initializer = True + +# Control what gets shown in documentation +breathe_show_include = False # Don't show #include directives + +# Order of member documentation +breathe_order_parameters_first = False + +# Use Sphinx's C++ domain for cross-references +breathe_use_project_refids = True + +# Suppress specific Breathe warnings for cleaner builds +breathe_debug_trace_directives = False +breathe_debug_trace_doxygen_ids = False +breathe_debug_trace_qualification = False intersphinx_mapping = {"pytorch": ("https://docs.pytorch.org/docs/main", None)} @@ -52,99 +92,6 @@ "misc.highlighting_failure", ] -# Configure Breathe -breathe_show_define_initializer = True -breathe_show_enumvalue_initializer = True -breathe_default_members = ("members", "undoc-members") - - -# Fix for Python 3.10+ compatibility with exhale 2.3.0 -# MutableMapping was moved from collections to collections.abc in Python 3.10 -try: - import collections - from collections.abc import MutableMapping - - if not hasattr(collections, "MutableMapping"): - collections.MutableMapping = MutableMapping -except ImportError: - pass - -# Setup absolute paths for communicating with breathe / exhale where -# items are expected / should be trimmed by. -# This file is {repo_root}/docs/cpp/source/conf.py -this_file_dir = os.path.abspath(os.path.dirname(__file__)) -doxygen_xml_dir = os.path.join( - os.path.dirname(this_file_dir), # {repo_root}/docs/cpp - "build", # {repo_root}/docs/cpp/build - "xml", # {repo_root}/docs/cpp/build/xml -) -repo_root = os.path.dirname( # {repo_root} - os.path.dirname( # {repo_root}/docs - os.path.dirname( # {repo_root}/docs/cpp - this_file_dir # {repo_root}/docs/cpp/source - ) - ) -) - -breathe_projects = {"PyTorch": doxygen_xml_dir} -breathe_default_project = "PyTorch" - -# Setup the exhale extension -exhale_args = { - ############################################################################ - # These arguments are required. # - ############################################################################ - "containmentFolder": "./api", - "rootFileName": "library_root.rst", - "rootFileTitle": "Library API", - "doxygenStripFromPath": repo_root, - ############################################################################ - # Suggested optional arguments. # - ############################################################################ - "createTreeView": True, - "exhaleExecutesDoxygen": True, - "exhaleUseDoxyfile": True, - "verboseBuild": True, - ############################################################################ - # HTML Theme specific configurations. # - ############################################################################ - # Fix broken Sphinx RTD Theme 'Edit on GitHub' links - # Search for 'Edit on GitHub' on the FAQ: - # http://exhale.readthedocs.io/en/latest/faq.html - "pageLevelConfigMeta": ":github_url: https://github.com/pytorch/pytorch", - ############################################################################ - # Individual page layout example configuration. # - ############################################################################ - # Example of adding contents directives on custom kinds with custom title - "contentsTitle": "Page Contents", - "kindsWithContentsDirectives": ["class", "file", "namespace", "struct"], - # Exclude PIMPL files from class hierarchy tree and namespace pages. - "listingExclude": [r".*Impl$"], - ############################################################################ - # Main library page layout example configuration. # - ############################################################################ - "afterTitleDescription": textwrap.dedent( - """ - Welcome to the developer reference for the PyTorch C++ API. - """ - ), - ############################################################################ - # Duplicate handling and error management. # - ############################################################################ - # Note: Using Doxyfile instead of stdin configuration - # "exhaleDoxygenStdin" is not compatible with "exhaleUseDoxyfile" - # Handle unresolved references more gracefully - "unabridgedOrphanKinds": { - "function", - "define", - "enum", - "enumvalue", - "typedef", - "variable", - }, - "fullToctreeMaxDepth": 2, -} - # Tell sphinx what the primary language being documented is. primary_domain = "cpp" @@ -164,7 +111,10 @@ # You can specify multiple suffix as a list of string: # # source_suffix = ['.rst', '.md'] -source_suffix = ".rst" +source_suffix = { + ".rst": "restructuredtext", + ".md": "markdown", +} # The master toctree document. master_doc = "index" @@ -264,8 +214,24 @@ # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". # NOTE: sharing python docs resources -html_static_path = [os.path.join(repo_root, "docs", "cpp", "source", "_static")] -html_css_files = ["cpp_theme.css"] + + +# Remove "Subclassed by" paragraphs that Breathe renders as plain text +# (not links) because per-class pages don't exist without exhale. +def remove_subclassed_by(app, doctree, docname): + from docutils import nodes + + for node in doctree.traverse(nodes.paragraph): + text = node.astext() + if text.startswith("Subclassed by "): + # Keep if it contains actual reference links + if not node.traverse(nodes.reference): + node.parent.remove(node) + + +def setup(app): + app.connect("doctree-resolved", remove_subclassed_by) + # Called automatically by Sphinx, making this `conf.py` an "extension". diff --git a/docs/cpp/source/faq.md b/docs/cpp/source/faq.md new file mode 100644 index 0000000000000..3ee305ccdf3ad --- /dev/null +++ b/docs/cpp/source/faq.md @@ -0,0 +1,244 @@ +--- +myst: + html_meta: + description: Frequently asked questions about the PyTorch C++ API and libtorch. + keywords: PyTorch, C++, FAQ, libtorch, troubleshooting +--- + +# FAQ + +Listed below are a number of common issues users face with the various parts of +the C++ API. + +## C++ Extensions + +### Undefined symbol errors from PyTorch/ATen + +**Problem**: You import your extension and get an `ImportError` stating that +some C++ symbol from PyTorch or ATen is undefined. For example: + +```cpp +>>> import extension +Traceback (most recent call last): + File "", line 1, in +ImportError: /home/user/.pyenv/versions/3.7.1/lib/python3.7/site-packages/extension.cpython-37m-x86_64-linux-gnu.so: undefined symbol: _ZN2at19UndefinedTensorImpl10_singletonE +``` + +**Fix**: The fix is to `import torch` before you import your extension. This will make +the symbols from the PyTorch dynamic (shared) library that your extension +depends on available, allowing them to be resolved once you import your extension. + +### I created a tensor using a function from `at::` and get errors + +**Problem**: You created a tensor using e.g. `at::ones` or `at::randn` or +any other tensor factory from the `at::` namespace and are getting errors. + +**Fix**: Replace `at::` with `torch::` for factory function calls. You +should never use factory functions from the `at::` namespace, as they will +create tensors. The corresponding `torch::` functions will create variables, +and you should only ever deal with variables in your code. + +## LibTorch + +### How do I move a model to GPU? + +**Problem**: You want to run your model on GPU but are unsure how to move both +the model and tensors to the correct device. + +**Fix**: Use the `to()` method to move your model and tensors to a CUDA device: + +```cpp +torch::Device device(torch::kCUDA); +model->to(device); +auto input = torch::randn({1, 3, 224, 224}).to(device); +auto output = model->forward(input); +``` + +You can also check for CUDA availability before moving: + +```cpp +torch::Device device(torch::cuda::is_available() ? torch::kCUDA : torch::kCPU); +``` + +Make sure to compile with the TorchScript headers by including ``. + +### My model runs slower in C++ than in Python + +**Problem**: Your model inference is slower in C++ compared to Python. + +**Fix**: There are several common causes: + +1. **Enable inference mode**: Wrap your inference code with `torch::NoGradGuard` + to disable gradient computation: + +```cpp +torch::NoGradGuard no_grad; +auto output = model->forward(input); +``` + +2. **Enable optimizations**: For TorchScript models, use `optimize_for_inference`: + +```cpp +module = torch::jit::optimize_for_inference(module); +``` + +3. **Warm up the model**: Run a few inference passes before benchmarking to allow + JIT compilation and memory allocation to complete. + +4. **Check thread settings**: Ensure proper thread configuration: + +```cpp +at::set_num_threads(4); // Adjust based on your hardware +``` + +## Neural Network Modules + +### How do I register submodules in a custom module? + +**Problem**: You created a custom module but the submodules are not being +recognized during `forward()` or when saving/loading the model. + +**Fix**: You must register submodules in the constructor using +`register_module()`: + +```cpp +struct MyModel : torch::nn::Module { + MyModel() { + fc1 = register_module("fc1", torch::nn::Linear(784, 128)); + fc2 = register_module("fc2", torch::nn::Linear(128, 10)); + } + + torch::Tensor forward(torch::Tensor x) { + x = torch::relu(fc1->forward(x)); + return fc2->forward(x); + } + + torch::nn::Linear fc1{nullptr}, fc2{nullptr}; +}; +``` + +### How do I set a module to evaluation mode? + +**Problem**: Layers like Dropout and BatchNorm behave differently during training +and evaluation, and you need to switch between modes. + +**Fix**: Use the `eval()` and `train()` methods: + +```cpp +model->eval(); // Set to evaluation mode +// ... run inference ... +model->train(); // Set back to training mode +``` + +## Data Loading + +### How do I create a custom dataset? + +**Problem**: You want to load your own data instead of using built-in datasets. + +**Fix**: Create a class that inherits from `torch::data::datasets::Dataset` and +implement the `get()` and `size()` methods: + +```cpp +class CustomDataset : public torch::data::datasets::Dataset { + public: + explicit CustomDataset(const std::string& data_path) { + // Load your data here + } + + torch::data::Example<> get(size_t index) override { + // Return a single data sample + torch::Tensor data = /* load data at index */; + torch::Tensor label = /* load label at index */; + return {data, label}; + } + + torch::optional size() const override { + return dataset_size_; + } + + private: + size_t dataset_size_; +}; +``` + +Then use it with a DataLoader: + +```cpp +auto dataset = CustomDataset("path/to/data") + .map(torch::data::transforms::Stack<>()); +auto dataloader = torch::data::make_data_loader( + std::move(dataset), + torch::data::DataLoaderOptions().batch_size(32).workers(4)); +``` + +## Serialization + +### How do I save and load model weights? + +**Problem**: You want to save trained model weights and load them later. + +**Fix**: Use `torch::save()` and `torch::load()`: + +```cpp +// Saving +torch::save(model, "model.pt"); + +// Loading +torch::load(model, "model.pt"); +``` + +For saving only specific tensors or state: + +```cpp +torch::serialize::OutputArchive archive; +model->save(archive); +archive.save_to("model_weights.pt"); + +// Loading +torch::serialize::InputArchive archive; +archive.load_from("model_weights.pt"); +model->load(archive); +``` + +## Build and Compilation + +### CMake cannot find Torch + +**Problem**: When building your project with CMake, you get an error that +`Torch` package cannot be found. + +**Fix**: You need to specify the path to the LibTorch installation using +`CMAKE_PREFIX_PATH`: + +```cpp +cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch .. +``` + +Alternatively, set `Torch_DIR` to point to the directory containing +`TorchConfig.cmake`: + +```cpp +cmake -DTorch_DIR=/path/to/libtorch/share/cmake/Torch .. +``` + +### Linker errors with undefined references + +**Problem**: Your project compiles but you get linker errors with undefined +references to PyTorch symbols. + +**Fix**: Ensure you're linking against all required libraries in your +`CMakeLists.txt`: + +```cpp +find_package(Torch REQUIRED) +add_executable(my_app main.cpp) +target_link_libraries(my_app "${TORCH_LIBRARIES}") +set_property(TARGET my_app PROPERTY CXX_STANDARD 17) +``` + +Also ensure that the compiler flags are set correctly: + +```cpp +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}") +``` diff --git a/docs/cpp/source/frontend.md b/docs/cpp/source/frontend.md new file mode 100644 index 0000000000000..244e2c836bf98 --- /dev/null +++ b/docs/cpp/source/frontend.md @@ -0,0 +1,146 @@ +--- +myst: + html_meta: + description: PyTorch C++ Frontend guide — defining models, training loops, and using torch::nn modules in C++. + keywords: PyTorch, C++, frontend, torch::nn, Module, training, inference +--- + +# The C++ Frontend + +The PyTorch C++ frontend is a C++17 library for CPU and GPU +tensor computation, with automatic differentiation and high level building +blocks for state of the art machine learning applications. + +## Description + +The PyTorch C++ frontend can be thought of as a C++ version of the +PyTorch Python frontend, providing automatic differentiation and various higher +level abstractions for machine learning and neural networks. Specifically, +it consists of the following components: + +| Component | Description | +| --- | --- | +| `torch::Tensor` | Automatically differentiable, efficient CPU and GPU enabled tensors | +| `torch::nn` | A collection of composable modules for neural network modeling | +| `torch::optim` | Optimization algorithms like SGD, Adam or RMSprop to train your models | +| `torch::data` | Datasets, data pipelines and multi-threaded, asynchronous data loader | +| `torch::serialize` | A serialization API for storing and loading model checkpoints | +| `torch::python` | Glue to bind your C++ models into Python | +| `torch::jit` | Pure C++ access to the TorchScript JIT compiler | + +## End-to-end example + +Here is a simple, end-to-end example of defining and training a simple +neural network on the MNIST dataset: + +```cpp +#include + +// Define a new Module. +struct Net : torch::nn::Module { + Net() { + // Construct and register two Linear submodules. + fc1 = register_module("fc1", torch::nn::Linear(784, 64)); + fc2 = register_module("fc2", torch::nn::Linear(64, 32)); + fc3 = register_module("fc3", torch::nn::Linear(32, 10)); + } + + // Implement the Net's algorithm. + torch::Tensor forward(torch::Tensor x) { + // Use one of many tensor manipulation functions. + x = torch::relu(fc1->forward(x.reshape({x.size(0), 784}))); + x = torch::dropout(x, /*p=*/0.5, /*train=*/is_training()); + x = torch::relu(fc2->forward(x)); + x = torch::log_softmax(fc3->forward(x), /*dim=*/1); + return x; + } + + // Use one of many "standard library" modules. + torch::nn::Linear fc1{nullptr}, fc2{nullptr}, fc3{nullptr}; +}; + +int main() { + // Create a new Net. + auto net = std::make_shared(); + + // Create a multi-threaded data loader for the MNIST dataset. + auto data_loader = torch::data::make_data_loader( + torch::data::datasets::MNIST("./data").map( + torch::data::transforms::Stack<>()), + /*batch_size=*/64); + + // Instantiate an SGD optimization algorithm to update our Net's parameters. + torch::optim::SGD optimizer(net->parameters(), /*lr=*/0.01); + + for (size_t epoch = 1; epoch <= 10; ++epoch) { + size_t batch_index = 0; + // Iterate the data loader to yield batches from the dataset. + for (auto& batch : *data_loader) { + // Reset gradients. + optimizer.zero_grad(); + // Execute the model on the input data. + torch::Tensor prediction = net->forward(batch.data); + // Compute a loss value to judge the prediction of our model. + torch::Tensor loss = torch::nll_loss(prediction, batch.target); + // Compute gradients of the loss w.r.t. the parameters of our model. + loss.backward(); + // Update the parameters based on the calculated gradients. + optimizer.step(); + // Output the loss and checkpoint every 100 batches. + if (++batch_index % 100 == 0) { + std::cout << "Epoch: " << epoch << " | Batch: " << batch_index + << " | Loss: " << loss.item() << std::endl; + // Serialize your model periodically as a checkpoint. + torch::save(net, "net.pt"); + } + } + } +} +``` + +To see more complete examples of using the PyTorch C++ frontend, see [the example repository](https://github.com/pytorch/examples/tree/master/cpp). + +## Philosophy + +PyTorch's C++ frontend was designed with the idea that the Python frontend is +great, and should be used when possible; but in some settings, performance and +portability requirements make the use of the Python interpreter infeasible. For +example, Python is a poor choice for low latency, high performance or +multithreaded environments, such as video games or production servers. The +goal of the C++ frontend is to address these use cases, while not sacrificing +the user experience of the Python frontend. + +As such, the C++ frontend has been written with a few philosophical goals in mind: + +* **Closely model the Python frontend in its design**, naming, conventions and + functionality. While there may be occasional differences between the two + frontends (e.g., where we have dropped deprecated features or fixed "warts" + in the Python frontend), we guarantee that the effort in porting a Python model + to C++ should lie exclusively in **translating language features**, + not modifying functionality or behavior. + +* **Prioritize flexibility and user-friendliness over micro-optimization.** + In C++, you can often get optimal code, but at the cost of an extremely + unfriendly user experience. Flexibility and dynamism is at the heart of + PyTorch, and the C++ frontend seeks to preserve this experience, in some + cases sacrificing performance (or "hiding" performance knobs) to keep APIs + simple and explicable. We want researchers who don't write C++ for a living + to be able to use our APIs. + +A word of warning: Python is not necessarily slower than +C++! The Python frontend calls into C++ for almost anything computationally expensive +(especially any kind of numeric operation), and these operations will take up +the bulk of time spent in a program. If you would prefer to write Python, +and can afford to write Python, we recommend using the Python interface to +PyTorch. However, if you would prefer to write C++, or need to write C++ +(because of multithreading, latency or deployment requirements), the +C++ frontend to PyTorch provides an API that is approximately as convenient, +flexible, friendly and intuitive as its Python counterpart. The two frontends +serve different use cases, work hand in hand, and neither is meant to +unconditionally replace the other. + +## Installation + +Instructions on how to install the C++ frontend library distribution, including +an example for how to build a minimal application depending on LibTorch, may be +found by following [this](https://pytorch.org/cppdocs/installing.html) link. diff --git a/docs/cpp/source/frontend.rst b/docs/cpp/source/frontend.rst deleted file mode 100644 index 7a1776f7bd4a6..0000000000000 --- a/docs/cpp/source/frontend.rst +++ /dev/null @@ -1,153 +0,0 @@ -The C++ Frontend -================ - -The PyTorch C++ frontend is a C++17 library for CPU and GPU -tensor computation, with automatic differentiation and high level building -blocks for state of the art machine learning applications. - -Description ------------ - -The PyTorch C++ frontend can be thought of as a C++ version of the -PyTorch Python frontend, providing automatic differentiation and various higher -level abstractions for machine learning and neural networks. Specifically, -it consists of the following components: - -+----------------------+------------------------------------------------------------------------+ -| Component | Description | -+======================+========================================================================+ -| ``torch::Tensor`` | Automatically differentiable, efficient CPU and GPU enabled tensors | -+----------------------+------------------------------------------------------------------------+ -| ``torch::nn`` | A collection of composable modules for neural network modeling | -+----------------------+------------------------------------------------------------------------+ -| ``torch::optim`` | Optimization algorithms like SGD, Adam or RMSprop to train your models | -+----------------------+------------------------------------------------------------------------+ -| ``torch::data`` | Datasets, data pipelines and multi-threaded, asynchronous data loader | -+----------------------+------------------------------------------------------------------------+ -| ``torch::serialize`` | A serialization API for storing and loading model checkpoints | -+----------------------+------------------------------------------------------------------------+ -| ``torch::python`` | Glue to bind your C++ models into Python | -+----------------------+------------------------------------------------------------------------+ -| ``torch::jit`` | Pure C++ access to the TorchScript JIT compiler | -+----------------------+------------------------------------------------------------------------+ - -End-to-end example ------------------- - -Here is a simple, end-to-end example of defining and training a simple -neural network on the MNIST dataset: - -.. code-block:: cpp - - #include - - // Define a new Module. - struct Net : torch::nn::Module { - Net() { - // Construct and register two Linear submodules. - fc1 = register_module("fc1", torch::nn::Linear(784, 64)); - fc2 = register_module("fc2", torch::nn::Linear(64, 32)); - fc3 = register_module("fc3", torch::nn::Linear(32, 10)); - } - - // Implement the Net's algorithm. - torch::Tensor forward(torch::Tensor x) { - // Use one of many tensor manipulation functions. - x = torch::relu(fc1->forward(x.reshape({x.size(0), 784}))); - x = torch::dropout(x, /*p=*/0.5, /*train=*/is_training()); - x = torch::relu(fc2->forward(x)); - x = torch::log_softmax(fc3->forward(x), /*dim=*/1); - return x; - } - - // Use one of many "standard library" modules. - torch::nn::Linear fc1{nullptr}, fc2{nullptr}, fc3{nullptr}; - }; - - int main() { - // Create a new Net. - auto net = std::make_shared(); - - // Create a multi-threaded data loader for the MNIST dataset. - auto data_loader = torch::data::make_data_loader( - torch::data::datasets::MNIST("./data").map( - torch::data::transforms::Stack<>()), - /*batch_size=*/64); - - // Instantiate an SGD optimization algorithm to update our Net's parameters. - torch::optim::SGD optimizer(net->parameters(), /*lr=*/0.01); - - for (size_t epoch = 1; epoch <= 10; ++epoch) { - size_t batch_index = 0; - // Iterate the data loader to yield batches from the dataset. - for (auto& batch : *data_loader) { - // Reset gradients. - optimizer.zero_grad(); - // Execute the model on the input data. - torch::Tensor prediction = net->forward(batch.data); - // Compute a loss value to judge the prediction of our model. - torch::Tensor loss = torch::nll_loss(prediction, batch.target); - // Compute gradients of the loss w.r.t. the parameters of our model. - loss.backward(); - // Update the parameters based on the calculated gradients. - optimizer.step(); - // Output the loss and checkpoint every 100 batches. - if (++batch_index % 100 == 0) { - std::cout << "Epoch: " << epoch << " | Batch: " << batch_index - << " | Loss: " << loss.item() << std::endl; - // Serialize your model periodically as a checkpoint. - torch::save(net, "net.pt"); - } - } - } - } - -To see more complete examples of using the PyTorch C++ frontend, see `the example repository -`_. - -Philosophy ----------- - -PyTorch's C++ frontend was designed with the idea that the Python frontend is -great, and should be used when possible; but in some settings, performance and -portability requirements make the use of the Python interpreter infeasible. For -example, Python is a poor choice for low latency, high performance or -multithreaded environments, such as video games or production servers. The -goal of the C++ frontend is to address these use cases, while not sacrificing -the user experience of the Python frontend. - -As such, the C++ frontend has been written with a few philosophical goals in mind: - -* **Closely model the Python frontend in its design**, naming, conventions and - functionality. While there may be occasional differences between the two - frontends (e.g., where we have dropped deprecated features or fixed "warts" - in the Python frontend), we guarantee that the effort in porting a Python model - to C++ should lie exclusively in **translating language features**, - not modifying functionality or behavior. - -* **Prioritize flexibility and user-friendliness over micro-optimization.** - In C++, you can often get optimal code, but at the cost of an extremely - unfriendly user experience. Flexibility and dynamism is at the heart of - PyTorch, and the C++ frontend seeks to preserve this experience, in some - cases sacrificing performance (or "hiding" performance knobs) to keep APIs - simple and explicable. We want researchers who don't write C++ for a living - to be able to use our APIs. - -A word of warning: Python is not necessarily slower than -C++! The Python frontend calls into C++ for almost anything computationally expensive -(especially any kind of numeric operation), and these operations will take up -the bulk of time spent in a program. If you would prefer to write Python, -and can afford to write Python, we recommend using the Python interface to -PyTorch. However, if you would prefer to write C++, or need to write C++ -(because of multithreading, latency or deployment requirements), the -C++ frontend to PyTorch provides an API that is approximately as convenient, -flexible, friendly and intuitive as its Python counterpart. The two frontends -serve different use cases, work hand in hand, and neither is meant to -unconditionally replace the other. - -Installation ------------- - -Instructions on how to install the C++ frontend library distribution, including -an example for how to build a minimal application depending on LibTorch, may be -found by following `this `_ link. diff --git a/docs/cpp/source/index.rst b/docs/cpp/source/index.md similarity index 54% rename from docs/cpp/source/index.rst rename to docs/cpp/source/index.md index fb7e5986e869a..1979721ad7a94 100644 --- a/docs/cpp/source/index.rst +++ b/docs/cpp/source/index.md @@ -1,5 +1,11 @@ -PyTorch C++ API -=============== +--- +myst: + html_meta: + description: PyTorch C++ API documentation — ATen tensors, Autograd, C++ Frontend, TorchScript, and C++ Extensions. + keywords: PyTorch, C++, API, libtorch, ATen, Autograd, TorchScript, C++ Frontend +--- + +# PyTorch C++ API These pages provide the documentation for the public portions of the PyTorch C++ API. This API can roughly be divided into five parts: @@ -16,108 +22,99 @@ networks with strong emphasis on GPU acceleration as well as fast CPU performance. It is currently in use at Facebook in research and production; we are looking forward to welcoming more users of the PyTorch C++ API. -.. warning:: +```{warning} - At the moment, the C++ API should be considered "beta" stability; we may - make major breaking changes to the backend in order to improve the API, - or in service of providing the Python interface to PyTorch, which is our - most stable and best supported interface. +At the moment, the C++ API should be considered "beta" stability; we may +make major breaking changes to the backend in order to improve the API, +or in service of providing the Python interface to PyTorch, which is our +most stable and best supported interface. +``` -ATen ----- +## ATen ATen is fundamentally a tensor library, on top of which almost all other Python -and C++ interfaces in PyTorch are built. It provides a core ``Tensor`` class, +and C++ interfaces in PyTorch are built. It provides a core `Tensor` class, on which many hundreds of operations are defined. Most of these operations have -both CPU and GPU implementations, to which the ``Tensor`` class will +both CPU and GPU implementations, to which the `Tensor` class will dynamically dispatch based on its type. A small example of using ATen could look as follows: -.. code-block:: cpp - - #include +```cpp +#include - at::Tensor a = at::ones({2, 2}, at::kInt); - at::Tensor b = at::randn({2, 2}); - auto c = a + b.to(at::kInt); +at::Tensor a = at::ones({2, 2}, at::kInt); +at::Tensor b = at::randn({2, 2}); +auto c = a + b.to(at::kInt); +``` -This ``Tensor`` class and all other symbols in ATen are found in the ``at::`` +This `Tensor` class and all other symbols in ATen are found in the `at::` namespace, documented -`here `_. +[here](https://pytorch.org/cppdocs/api/namespace_at.html#namespace-at). -Autograd --------- +## Autograd What we term *autograd* are the portions of PyTorch's C++ API that augment the -ATen ``Tensor`` class with capabilities concerning automatic differentiation. +ATen `Tensor` class with capabilities concerning automatic differentiation. The autograd system records operations on tensors to form an *autograd graph*. -Calling ``backwards()`` on a leaf variable in this graph performs reverse mode +Calling `backwards()` on a leaf variable in this graph performs reverse mode differentiation through the network of functions and tensors spanning the autograd graph, ultimately yielding gradients. The following example provides a taste of this interface: -.. code-block:: cpp - - #include - #include +```cpp +#include +#include - torch::Tensor a = torch::ones({2, 2}, torch::requires_grad()); - torch::Tensor b = torch::randn({2, 2}); - auto c = a + b; - c.backward(); // a.grad() will now hold the gradient of c w.r.t. a. +torch::Tensor a = torch::ones({2, 2}, torch::requires_grad()); +torch::Tensor b = torch::randn({2, 2}); +auto c = a + b; +c.backward(); // a.grad() will now hold the gradient of c w.r.t. a. +``` -The ``at::Tensor`` class in ATen is not differentiable by default. To add the +The `at::Tensor` class in ATen is not differentiable by default. To add the differentiability of tensors the autograd API provides, you must use tensor factory functions from the `torch::` namespace instead of the `at::` namespace. For example, while a tensor created with `at::ones` will not be differentiable, a tensor created with `torch::ones` will be. -C++ Frontend ------------- +## C++ Frontend The PyTorch C++ frontend provides a high level, pure C++ modeling interface for neural networks and general ML (Machine Learning) research and production use cases, largely following the Python API in design and provided functionality. The C++ frontend includes the following: -- An interface for defining machine learning models through a hierarchical module system (like ``torch.nn.Module``); +- An interface for defining machine learning models through a hierarchical module system (like `torch.nn.Module`); - A "standard library" of pre-existing modules for the most common modeling purposes (e.g. convolutions, RNNs, batch normalization etc.); - An optimization API, including implementations of popular optimizers such as SGD, Adam, RMSprop and others; - A means of representing datasets and data pipelines, including functionality to load data in parallel over many CPU cores; -- A serialization format for storing and loading checkpoints of a training session (like ``torch.utils.data.DataLoader``); -- Automatic parallelization of models onto multiple GPUs (like ``torch.nn.parallel.DataParallel``); +- A serialization format for storing and loading checkpoints of a training session (like `torch.utils.data.DataLoader`); +- Automatic parallelization of models onto multiple GPUs (like `torch.nn.parallel.DataParallel`); - Support code to easily bind C++ models into Python using pybind11; - Entry points to the TorchScript JIT compiler; - Helpful utilities to facilitate interfacing with the ATen and Autograd APIs. -See `this document `_ for a more +See [this document](https://pytorch.org/cppdocs/frontend.html) for a more detailed description of the C++ frontend. Relevant sections of the `torch::` -namespace related to the C++ Frontend include `torch::nn -`_, -`torch::optim -`_, -`torch::data -`_, -`torch::serialize -`_, -`torch::jit -`_ -and `torch::python -`_. -Examples of the C++ frontend can be found in `this repository -`_ which is being +namespace related to the C++ Frontend include [torch::nn](https://pytorch.org/cppdocs/api/namespace_torch__nn.html#namespace-torch-nn), +[torch::optim](https://pytorch.org/cppdocs/api/namespace_torch__optim.html#namespace-torch-optim), +[torch::data](https://pytorch.org/cppdocs/api/namespace_torch__data.html#namespace-torch-data), +[torch::serialize](https://pytorch.org/cppdocs/api/namespace_torch__serialize.html#namespace-torch-serialize), +[torch::jit](https://pytorch.org/cppdocs/api/namespace_torch__jit.html#namespace-torch-jit) +and [torch::python](https://pytorch.org/cppdocs/api/namespace_torch__python.html#namespace-torch-python). +Examples of the C++ frontend can be found in [this repository](https://github.com/pytorch/examples/tree/master/cpp) which is being expanded on a continuous and active basis. -.. note:: +```{note} - Unless you have a particular reason to constrain yourself exclusively to ATen - or the Autograd API, the C++ frontend is the recommended entry point to the - PyTorch C++ ecosystem. While it is still in beta as we collect user feedback - (from you!), it provides both more functionality and better stability - guarantees than the ATen and Autograd APIs. +Unless you have a particular reason to constrain yourself exclusively to ATen +or the Autograd API, the C++ frontend is the recommended entry point to the +PyTorch C++ ecosystem. While it is still in beta as we collect user feedback +(from you!), it provides both more functionality and better stability +guarantees than the ATen and Autograd APIs. +``` -TorchScript ------------ +## TorchScript TorchScript is a representation of a PyTorch model that can be understood, compiled and serialized by the TorchScript compiler. Fundamentally, TorchScript @@ -132,16 +129,13 @@ functionality: The first mechanism may be of great interest to you if you would like to define your models in Python as much as possible, but subsequently export them to C++ for production environments and no-Python inference. You can find out more -about this by following `this -`_ link. The second +about this by following [this](https://pytorch.org/tutorials/advanced/cpp_export.html) link. The second API concerns itself with scenarios in which you would like to extend TorchScript with custom operators, which can similarly be serialized and -invoked from C++ during inference. Lastly, the `torch::jit::compile -`_ +invoked from C++ during inference. Lastly, the [torch::jit::compile](https://pytorch.org/cppdocs/api/function_namespacetorch_1_1jit_1a8660dc13a6b82336aadac667e6dccba1.html) function may be used to access the TorchScript compiler directly from C++. -C++ Extensions --------------- +## C++ Extensions *C++ Extensions* offer a simple yet powerful way of accessing all of the above interfaces for the purpose of extending regular Python use-cases of PyTorch. @@ -151,37 +145,26 @@ does not add any new functionality to the PyTorch C++ API. Instead, it provides integration with Python setuptools as well as JIT compilation mechanisms that allow access to ATen, the autograd and other C++ APIs from Python. To learn more about the C++ extension API, go through -`this tutorial `_. - -Contents --------- - -.. toctree:: - :maxdepth: 2 +[this tutorial](https://pytorch.org/tutorials/advanced/cpp_extension.html). - installing - frontend - stable - api/library_root +## Contents -.. toctree:: - :glob: - :maxdepth: 1 - :caption: Notes +```{toctree} +:maxdepth: 2 - notes/* +installing +frontend +api/index +faq +``` -Indices and tables -================== +# Indices and tables -* :ref:`genindex` -* :ref:`modindex` -* :ref:`search` +* {ref}`genindex` +* {ref}`modindex` +* {ref}`search` -Acknowledgements ----------------- +## Acknowledgements -This documentation website for the PyTorch C++ universe has been enabled by the -`Exhale `_ project and generous investment -of time and effort by its maintainer, `svenevs `_. -We thank Stephen for his work and his efforts providing help with the PyTorch C++ documentation. +This documentation website for the PyTorch C++ universe uses the Sphinx +C++ domain for API documentation. diff --git a/docs/cpp/source/installing.md b/docs/cpp/source/installing.md new file mode 100644 index 0000000000000..7fb9067b9611f --- /dev/null +++ b/docs/cpp/source/installing.md @@ -0,0 +1,187 @@ +--- +myst: + html_meta: + description: How to install and configure the PyTorch C++ API (libtorch) for development. + keywords: PyTorch, C++, install, libtorch, setup, CMake, build +--- + +# Installing C++ Distributions of PyTorch + +We provide binary distributions of all headers, libraries and CMake +configuration files required to depend on PyTorch. We call this distribution +*LibTorch*, and you can download ZIP archives containing the latest LibTorch +distribution on [our website](https://pytorch.org/get-started/locally/). Below +is a small example of writing a minimal application that depends on LibTorch +and uses the `torch::Tensor` class which comes with the PyTorch C++ API. + +## Minimal Example + +The first step is to download the LibTorch ZIP archive via the link above. For +example: + +```sh +wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip +unzip libtorch-shared-with-deps-latest.zip +``` + +Note that the above link has CPU-only libtorch. If you would like to download a GPU-enabled +libtorch, find the right link in the link selector on https://pytorch.org + +If you're a Windows developer and wouldn't like to use CMake, you could jump to the Visual Studio +Extension section. + +Next, we can write a minimal CMake build configuration to develop a small +application that depends on LibTorch. CMake is not a hard requirement for using +LibTorch, but it is the recommended and blessed build system and will be well +supported into the future. A most basic `CMakeLists.txt` file could look like +this: + +```cmake +cmake_minimum_required(VERSION 3.18 FATAL_ERROR) +project(example-app) + +find_package(Torch REQUIRED) +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}") + +add_executable(example-app example-app.cpp) +target_link_libraries(example-app "${TORCH_LIBRARIES}") +set_property(TARGET example-app PROPERTY CXX_STANDARD 20) + +# The following code block is suggested to be used on Windows. +# According to https://github.com/pytorch/pytorch/issues/25457, +# the DLLs need to be copied to avoid memory errors. +if (MSVC) + file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll") + add_custom_command(TARGET example-app + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different + ${TORCH_DLLS} + $) +endif (MSVC) +``` + +The implementation of our example will simply create a new `torch::Tensor` and +print it: + +```cpp +#include +#include + +int main() { + torch::Tensor tensor = torch::rand({2, 3}); + std::cout << tensor << std::endl; +} +``` + +While there are more fine-grained headers you can include to access only parts +of the PyTorch C++ API, including `torch/torch.h` is the most sure-proof way of +including most of its functionality. + +The last step is to build the application. For this, assume our example +directory is laid out like this: + +```sh +example-app/ + CMakeLists.txt + example-app.cpp +``` + +We can now run the following commands to build the application from within the +`example-app/` folder: + +```sh +mkdir build +cd build +cmake -DCMAKE_PREFIX_PATH=/absolute/path/to/libtorch .. +cmake --build . --config Release +``` + +where `/absolute/path/to/libtorch` should be the absolute (!) path to the unzipped LibTorch +distribution. If PyTorch was installed via pip, `CMAKE_PREFIX_PATH` can be queried +using `torch.utils.cmake_prefix_path` variable. In that case CMake configuration step would look something like follows: + +```sh +cmake -DCMAKE_PREFIX_PATH=`python3 -c 'import torch;print(torch.utils.cmake_prefix_path)'` .. +``` + +If all goes well, it will look something like this: + +```sh +root@4b5a67132e81:/example-app# mkdir build +root@4b5a67132e81:/example-app# cd build +root@4b5a67132e81:/example-app/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch .. +-- The C compiler identification is GNU 5.4.0 +-- The CXX compiler identification is GNU 5.4.0 +-- Check for working C compiler: /usr/bin/cc +-- Check for working C compiler: /usr/bin/cc -- works +-- Detecting C compiler ABI info +-- Detecting C compiler ABI info - done +-- Detecting C compile features +-- Detecting C compile features - done +-- Check for working CXX compiler: /usr/bin/c++ +-- Check for working CXX compiler: /usr/bin/c++ -- works +-- Detecting CXX compiler ABI info +-- Detecting CXX compiler ABI info - done +-- Detecting CXX compile features +-- Detecting CXX compile features - done +-- Looking for pthread.h +-- Looking for pthread.h - found +-- Looking for pthread_create +-- Looking for pthread_create - not found +-- Looking for pthread_create in pthreads +-- Looking for pthread_create in pthreads - not found +-- Looking for pthread_create in pthread +-- Looking for pthread_create in pthread - found +-- Found Threads: TRUE +-- Configuring done +-- Generating done +-- Build files have been written to: /example-app/build +root@4b5a67132e81:/example-app/build# cmake --build . --config Release +Scanning dependencies of target example-app +[ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o +[100%] Linking CXX executable example-app +[100%] Built target example-app +``` + +Executing the resulting `example-app` binary found in the `build` folder +should now merrily print the tensor (exact output subject to randomness): + +```sh +root@4b5a67132e81:/example-app/build# ./example-app +0.2063 0.6593 0.0866 +0.0796 0.5841 0.1569 +[ Variable[CPUFloatType]{2,3} ] +``` + +```{tip} + +On Windows, debug and release builds are not ABI-compatible. If you plan to +build your project in debug mode, please try the debug version of LibTorch. +Also, make sure you specify the correct configuration in the `cmake --build .` +line above. +``` + +## System Requirements + +To ensure smooth installation and usage of LibTorch, please ensure your system +meets the following requirements: + +1. **GLIBC Version**: + + - GLIBC 2.29 or newer for cxx11 ABI version + +2. **GCC Version**: + + - GCC 9 or newer for cxx11 + +## Visual Studio Extension + +[LibTorch Project Template](https://marketplace.visualstudio.com/items?itemName=YiZhang.LibTorch001) can help Windows developers +set all libtorch project settings and link options for debug and release. +It's easy to use and you could check out the [demo video](https://ossci-windows.s3.us-east-1.amazonaws.com/vsextension/demo.mp4). +The only prerequisite is to download the libtorch on https://pytorch.org + +## Support + +If you run into any troubles with this installation and minimal usage guide, +please use our [forum](https://discuss.pytorch.org/) or [GitHub issues](https://github.com/pytorch/pytorch/issues) to get in touch. diff --git a/docs/cpp/source/installing.rst b/docs/cpp/source/installing.rst deleted file mode 100644 index 0612d2592f21a..0000000000000 --- a/docs/cpp/source/installing.rst +++ /dev/null @@ -1,182 +0,0 @@ -Installing C++ Distributions of PyTorch -======================================= - -We provide binary distributions of all headers, libraries and CMake -configuration files required to depend on PyTorch. We call this distribution -*LibTorch*, and you can download ZIP archives containing the latest LibTorch -distribution on `our website `_. Below -is a small example of writing a minimal application that depends on LibTorch -and uses the ``torch::Tensor`` class which comes with the PyTorch C++ API. - -Minimal Example ---------------- - -The first step is to download the LibTorch ZIP archive via the link above. For -example: - -.. code-block:: sh - - wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-shared-with-deps-latest.zip - unzip libtorch-shared-with-deps-latest.zip - -Note that the above link has CPU-only libtorch. If you would like to download a GPU-enabled -libtorch, find the right link in the link selector on https://pytorch.org - -If you're a Windows developer and wouldn't like to use CMake, you could jump to the Visual Studio -Extension section. - -Next, we can write a minimal CMake build configuration to develop a small -application that depends on LibTorch. CMake is not a hard requirement for using -LibTorch, but it is the recommended and blessed build system and will be well -supported into the future. A most basic `CMakeLists.txt` file could look like -this: - -.. code-block:: cmake - - cmake_minimum_required(VERSION 3.18 FATAL_ERROR) - project(example-app) - - find_package(Torch REQUIRED) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}") - - add_executable(example-app example-app.cpp) - target_link_libraries(example-app "${TORCH_LIBRARIES}") - set_property(TARGET example-app PROPERTY CXX_STANDARD 17) - - # The following code block is suggested to be used on Windows. - # According to https://github.com/pytorch/pytorch/issues/25457, - # the DLLs need to be copied to avoid memory errors. - if (MSVC) - file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll") - add_custom_command(TARGET example-app - POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - ${TORCH_DLLS} - $) - endif (MSVC) - -The implementation of our example will simply create a new `torch::Tensor` and -print it: - -.. code-block:: cpp - - #include - #include - - int main() { - torch::Tensor tensor = torch::rand({2, 3}); - std::cout << tensor << std::endl; - } - -While there are more fine-grained headers you can include to access only parts -of the PyTorch C++ API, including `torch/torch.h` is the most sure-proof way of -including most of its functionality. - -The last step is to build the application. For this, assume our example -directory is laid out like this: - -.. code-block:: sh - - example-app/ - CMakeLists.txt - example-app.cpp - -We can now run the following commands to build the application from within the -``example-app/`` folder: - -.. code-block:: sh - - mkdir build - cd build - cmake -DCMAKE_PREFIX_PATH=/absolute/path/to/libtorch .. - cmake --build . --config Release - -where ``/absolute/path/to/libtorch`` should be the absolute (!) path to the unzipped LibTorch -distribution. If PyTorch was installed via pip, `CMAKE_PREFIX_PATH` can be queried -using `torch.utils.cmake_prefix_path` variable. In that case CMake configuration step would look something like follows: - -.. code-block:: sh - - cmake -DCMAKE_PREFIX_PATH=`python3 -c 'import torch;print(torch.utils.cmake_prefix_path)'` .. - -If all goes well, it will look something like this: - -.. code-block:: sh - - root@4b5a67132e81:/example-app# mkdir build - root@4b5a67132e81:/example-app# cd build - root@4b5a67132e81:/example-app/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch .. - -- The C compiler identification is GNU 5.4.0 - -- The CXX compiler identification is GNU 5.4.0 - -- Check for working C compiler: /usr/bin/cc - -- Check for working C compiler: /usr/bin/cc -- works - -- Detecting C compiler ABI info - -- Detecting C compiler ABI info - done - -- Detecting C compile features - -- Detecting C compile features - done - -- Check for working CXX compiler: /usr/bin/c++ - -- Check for working CXX compiler: /usr/bin/c++ -- works - -- Detecting CXX compiler ABI info - -- Detecting CXX compiler ABI info - done - -- Detecting CXX compile features - -- Detecting CXX compile features - done - -- Looking for pthread.h - -- Looking for pthread.h - found - -- Looking for pthread_create - -- Looking for pthread_create - not found - -- Looking for pthread_create in pthreads - -- Looking for pthread_create in pthreads - not found - -- Looking for pthread_create in pthread - -- Looking for pthread_create in pthread - found - -- Found Threads: TRUE - -- Configuring done - -- Generating done - -- Build files have been written to: /example-app/build - root@4b5a67132e81:/example-app/build# cmake --build . --config Release - Scanning dependencies of target example-app - [ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o - [100%] Linking CXX executable example-app - [100%] Built target example-app - -Executing the resulting ``example-app`` binary found in the ``build`` folder -should now merrily print the tensor (exact output subject to randomness): - -.. code-block:: sh - - root@4b5a67132e81:/example-app/build# ./example-app - 0.2063 0.6593 0.0866 - 0.0796 0.5841 0.1569 - [ Variable[CPUFloatType]{2,3} ] - -.. tip:: - On Windows, debug and release builds are not ABI-compatible. If you plan to - build your project in debug mode, please try the debug version of LibTorch. - Also, make sure you specify the correct configuration in the ``cmake --build .`` - line above. - -System Requirements -------------------- - -To ensure smooth installation and usage of LibTorch, please ensure your system -meets the following requirements: - -1. **GLIBC Version**: - - GLIBC 2.29 or newer for cxx11 ABI version - -2. **GCC Version**: - - GCC 9 or newer for cxx11 - -Visual Studio Extension ------------------------ - -`LibTorch Project Template `_ can help Windows developers -set all libtorch project settings and link options for debug and release. -It's easy to use and you could check out the `demo video `_. -The only prerequisite is to download the libtorch on https://pytorch.org - -Support -------- - -If you run into any troubles with this installation and minimal usage guide, -please use our `forum `_ or `GitHub issues -`_ to get in touch. diff --git a/docs/cpp/source/library.rst b/docs/cpp/source/library.rst deleted file mode 100644 index 5cbaaf959910b..0000000000000 --- a/docs/cpp/source/library.rst +++ /dev/null @@ -1,37 +0,0 @@ -Torch Library API -================= - -The PyTorch C++ API provides capabilities for extending PyTorch's core library -of operators with user defined operators and data types. Extensions implemented -using the Torch Library API are made available for use in both the PyTorch eager -API as well as in TorchScript. - -For a tutorial style introduction to the library API, check out the -`Extending TorchScript with Custom C++ Operators -`_ -tutorial. - -Macros ------- - -.. doxygendefine:: TORCH_LIBRARY - -.. doxygendefine:: TORCH_LIBRARY_IMPL - -Classes -------- - -.. doxygenclass:: torch::Library - :members: - -.. doxygenclass:: torch::CppFunction - :members: - -Functions ---------- - -.. doxygengroup:: torch-dispatch-overloads - :content-only: - -.. doxygengroup:: torch-schema-overloads - :content-only: diff --git a/docs/cpp/source/notes/faq.rst b/docs/cpp/source/notes/faq.rst deleted file mode 100644 index 37a1a609f3a5f..0000000000000 --- a/docs/cpp/source/notes/faq.rst +++ /dev/null @@ -1,34 +0,0 @@ -FAQ ---- - -Listed below are a number of common issues users face with the various parts of -the C++ API. - -C++ Extensions -============== - -Undefined symbol errors from PyTorch/ATen -***************************************** - -**Problem**: You import your extension and get an ``ImportError`` stating that -some C++ symbol from PyTorch or ATen is undefined. For example:: - - >>> import extension - Traceback (most recent call last): - File "", line 1, in - ImportError: /home/user/.pyenv/versions/3.7.1/lib/python3.7/site-packages/extension.cpython-37m-x86_64-linux-gnu.so: undefined symbol: _ZN2at19UndefinedTensorImpl10_singletonE - -**Fix**: The fix is to ``import torch`` before you import your extension. This will make -the symbols from the PyTorch dynamic (shared) library that your extension -depends on available, allowing them to be resolved once you import your extension. - -I created a tensor using a function from ``at::`` and get errors -**************************************************************** - -**Problem**: You created a tensor using e.g. ``at::ones`` or ``at::randn`` or -any other tensor factory from the ``at::`` namespace and are getting errors. - -**Fix**: Replace ``at::`` with ``torch::`` for factory function calls. You -should never use factory functions from the ``at::`` namespace, as they will -create tensors. The corresponding ``torch::`` functions will create variables, -and you should only ever deal with variables in your code. diff --git a/docs/cpp/source/notes/inference_mode.rst b/docs/cpp/source/notes/inference_mode.rst deleted file mode 100644 index 60cc48ca93d5b..0000000000000 --- a/docs/cpp/source/notes/inference_mode.rst +++ /dev/null @@ -1,121 +0,0 @@ -Inference Mode -============== - -``c10::InferenceMode`` is a new RAII guard analogous to ``NoGradMode`` -to be used when you are certain your operations will have no interactions -with autograd (e.g. model training). Compared to ``NoGradMode``, code run -under this mode gets better performance by disabling autograd related work like -view tracking and version counter bumps. However, tensors created inside -``c10::InferenceMode`` have more limitations when interacting with autograd system as well. - -``InferenceMode`` can be enabled for a given block of code. Inside ``InferenceMode`` -all newly allocated (non-view) tensors are marked as inference tensors. Inference tensors: - -- do not have a version counter so an error will be raised if you try to read their version - (e.g., because you saved this tensor for backward). -- are immutable outside ``InferenceMode``. So an error will be raised if you try to: - - mutate their data outside InferenceMode. - - mutate them into ``requires_grad=True`` outside InferenceMode. - To work around you can make a clone outside ``InferenceMode`` to get a normal tensor before mutating. - -A non-view tensor is an inference tensor if and only if it was allocated inside ``InferenceMode``. -A view tensor is an inference tensor if and only if it is a view of an inference tensor. - -Inside an ``InferenceMode`` block, we make the following performance guarantees: - -- Like ``NoGradMode``, all operations do not record ``grad_fn`` even if their inputs have ``requires_grad=True``. - This applies to both inference tensors and normal tensors. -- View operations on inference tensors do not do view tracking. View and non-view inference tensors are - indistinguishable. -- Inplace operations on inference tensors are guaranteed not to do a version bump. - -For more implementation details of ``InferenceMode`` please see the `RFC-0011-InferenceMode `_. - -Migration guide from ``AutoNonVariableTypeMode`` ------------------------------------------------- - -In production use of PyTorch for inference workload, we have seen a proliferation -of uses of the C++ guard ``AutoNonVariableTypeMode`` (now ``AutoDispatchBelowADInplaceOrView``), -which disables autograd, view tracking and version counter bumps. Unfortunately, -current colloquial of this guard for inference workload is unsafe: it's possible to -use ``AutoNonVariableTypeMode`` to bypass PyTorch's safety checks and result in -silently wrong results, e.g. PyTorch throws an error when tensors saved for backwards -are subsequently mutated, but mutation happens inside ``AutoNonVariableTypeMode`` will -silently bypass the check and returns wrong gradient to users. - -When current users of ``AutoNonVariableTypeMode`` think about migrating, the following -steps might help you decide the best alternatives: - -1. Users trying to run workload in inference only mode (like loading a pretrained JIT model and - run inference in C++ runtime) should add ``c10::InferenceMode guard`` to guard all operations - on tensors (including model loading). See an inference workload example below: - -.. code-block:: cpp - - c10::InferenceMode guard; - model.load_jit(saved_model); - auto inputs = preprocess_tensors(data); - auto out = model.forward(inputs); - auto outputs = postprocess_tensors(out); - -Note ``c10::InferenceMode`` offers a drop in replacement for ``AutoNonVariableTypeMode`` which preserves -the performance characteristics of ``AutoNonVariableTypeMode``. But they also have some differences that -users should pay additional attention to: - - - Both guards affects tensor execution process to skip work not related to inference, but ``InferenceMode`` - also affects tensor creation while ``AutoNonVariableTypeMode`` doesn't. In other words, tensors created - inside ``InferenceMode`` are marked as inference tensors so that certain limitations can be applied after - exiting ``InferenceMode``. - - Enabled/disabled ``InferenceMode`` states can be nested while ``AutoNonVariableTypeMode`` only allows enabled state. - -.. code-block:: cpp - - { - InferenceMode guard(true); - // InferenceMode is on - { - InferenceMode guard(false); - // InferenceMode is off - } - // InferenceMode is on - } - // InferenceMode is off - - -2. Users trying to implement a customized kernel who want to redispatch under ``Autograd`` dispatch - keys should use ``AutoDispatchBelowADInplaceOrView`` instead. Note ``AutoDispatchBelowADInplaceOrView`` is just a new name - of ``AutoNonVariableTypeMode`` since it explains the guard's functionality better. We're deprecating - ``AutoNonVariableTypeMode`` and it'll be removed in 1.10 release. See customized kernel - ``ROIAlignFunction`` in ``pytorch/vision`` for an example: - -.. code-block:: cpp - - class ROIAlignFunction : public torch::autograd::Function { - public: - static torch::autograd::variable_list forward( - torch::autograd::AutogradContext* ctx, - const torch::autograd::Variable& input, - const torch::autograd::Variable& rois, - double spatial_scale, - int64_t pooled_height, - int64_t pooled_width, - int64_t sampling_ratio, - bool aligned) { - ctx->saved_data["spatial_scale"] = spatial_scale; - ctx->saved_data["pooled_height"] = pooled_height; - ctx->saved_data["pooled_width"] = pooled_width; - ctx->saved_data["sampling_ratio"] = sampling_ratio; - ctx->saved_data["aligned"] = aligned; - ctx->saved_data["input_shape"] = input.sizes(); - ctx->save_for_backward({rois}); - // Used to be at::AutoNonVariableTypeMode g; - at::AutoDispatchBelowADInplaceOrView guard; - auto result = roi_align( - input, rois, spatial_scale, pooled_height, - pooled_width, sampling_ratio, aligned); - return {result}; - } - -Customized inplace & view kernels need some special handling in addition to the guard above, see -`custom kernel tutorial `_ -for more details. diff --git a/docs/cpp/source/notes/maybe_owned.rst b/docs/cpp/source/notes/maybe_owned.rst deleted file mode 100644 index a1bc2d02b9882..0000000000000 --- a/docs/cpp/source/notes/maybe_owned.rst +++ /dev/null @@ -1,59 +0,0 @@ -MaybeOwned -================== - -``MaybeOwned`` is a C++ smart pointer class that dynamically -encodes whether a Tensor is *owned* or *borrowed*. It is used in -certain performance-sensitive situations to avoid unnecessarily -incrementing a Tensor’s reference count (at a small cost in -overhead from the extra indirection). - -.. warning:: - MaybeOwned must be used with **extreme** care. Claims of (non-)ownership - are not statically checked, and mistakes can cause reference undercounting - and use-after-free crashes. - - Due to this lack of safety net, we discourage the use of MaybeOwned - outside code paths that are known to be highly performance sensitive. - However, if you encounter pre-existing uses of MaybeOwned in code that - you want to modify, it’s critical to understand how to use it correctly. - -The primary use case for ``MaybeOwned`` is a function or method that -dynamically chooses between returning one of its arguments (typically -from a passthrough or “no-op” code path) and returning a freshly constructed -Tensor. Such a function would return a ``MaybeOwned`` in both cases, -the former in a "borrowed" state via a call to ``MaybeOwned::borrowed()``, -and the latter in an "owned" state via a call to ``MaybeOwned::owned()``. - -The canonical example is ``Tensor``'s ``expect_contiguous`` method, which shortcuts -and returns a borrowed self-reference when already contiguous: - -.. code-block:: cpp - - inline c10::MaybeOwned Tensor::expect_contiguous(MemoryFormat memory_format) const & { - if (is_contiguous(memory_format)) { - return c10::MaybeOwned::borrowed(*this); - } else { - return c10::MaybeOwned::owned(__dispatch_contiguous(memory_format)); - } - } - -Using the vocabulary of lifetimes, the essential safety requirement for borrowing -is that a borrowed Tensor must outlive any borrowing references to it. Here, for -example, we can safely borrow ``*this``, but the Tensor returned by -``__dispatch_contiguous()`` is freshly created, and borrowing a reference would -effectively leave it ownerless. - -So, general rules of thumb: - -- When in doubt, don’t use ``MaybeOwned`` at all - in particular, prefer - avoiding using it in code that doesn’t use it already. New usage should only be - introduced when critical (and demonstrable) performance gains result. - -- When modifying or calling code that already uses ``MaybeOwned``, remember - that it's always safe to produce a ``MaybeOwned`` from a Tensor in hand - via a call to ``MaybeOwned::owned()``. This may result in an unnecessary - reference count, but never in misbehavior - so it's always the safer bet, unless - the lifetime of the Tensor you're looking to wrap is crystal clear. - -More details and implementation code can be found at and -. diff --git a/docs/cpp/source/notes/tensor_basics.rst b/docs/cpp/source/notes/tensor_basics.rst deleted file mode 100644 index cf8f68a24eec7..0000000000000 --- a/docs/cpp/source/notes/tensor_basics.rst +++ /dev/null @@ -1,154 +0,0 @@ -Tensor Basics -============= - -The ATen tensor library backing PyTorch is a simple tensor library that exposes -the Tensor operations in Torch directly in C++17. ATen's API is auto-generated -from the same declarations PyTorch uses so the two APIs will track each other -over time. - -Tensor types are resolved dynamically, such that the API is generic and does not -include templates. That is, there is one ``Tensor`` type. It can hold a CPU or -CUDA Tensor, and the tensor may have Doubles, Float, Ints, etc. This design -makes it easy to write generic code without templating everything. - -See https://pytorch.org/cppdocs/api/namespace_at.html#functions for the provided -API. Excerpt: - -.. code-block:: cpp - - Tensor atan2(const Tensor & other) const; - Tensor & atan2_(const Tensor & other); - Tensor pow(Scalar exponent) const; - Tensor pow(const Tensor & exponent) const; - Tensor & pow_(Scalar exponent); - Tensor & pow_(const Tensor & exponent); - Tensor lerp(const Tensor & end, Scalar weight) const; - Tensor & lerp_(const Tensor & end, Scalar weight); - Tensor histc() const; - Tensor histc(int64_t bins) const; - Tensor histc(int64_t bins, Scalar min) const; - Tensor histc(int64_t bins, Scalar min, Scalar max) const; - -In place operations are also provided, and always suffixed by `_` to indicate -they will modify the Tensor. - -Efficient Access to Tensor Elements ------------------------------------ - -When using Tensor-wide operations, the relative cost of dynamic dispatch is very -small. However, there are cases, especially in your own kernels, where efficient -element-wise access is needed, and the cost of dynamic dispatch inside the -element-wise loop is very high. ATen provides *accessors* that are created with -a single dynamic check that a Tensor is the type and number of dimensions. -Accessors then expose an API for accessing the Tensor elements efficiently. - -Accessors are temporary views of a Tensor. They are only valid for the lifetime -of the tensor that they view and hence should only be used locally in a -function, like iterators. - -Note that accessors are not compatible with CUDA tensors inside kernel functions. -Instead, you will have to use a *packed accessor* which behaves the same way but -copies tensor metadata instead of pointing to it. - -It is thus recommended to use *accessors* for CPU tensors and *packed accessors* -for CUDA tensors. - -CPU accessors -************* - -.. code-block:: cpp - - torch::Tensor foo = torch::rand({12, 12}); - - // assert foo is 2-dimensional and holds floats. - auto foo_a = foo.accessor(); - float trace = 0; - - for(int i = 0; i < foo_a.size(0); i++) { - // use the accessor foo_a to get tensor data. - trace += foo_a[i][i]; - } - -CUDA accessors -************** - - -.. code-block:: cpp - - __global__ void packed_accessor_kernel( - torch::PackedTensorAccessor64 foo, - float* trace) { - int i = threadIdx.x; - gpuAtomicAdd(trace, foo[i][i]); - } - - torch::Tensor foo = torch::rand({12, 12}); - - // assert foo is 2-dimensional and holds floats. - auto foo_a = foo.packed_accessor64(); - float trace = 0; - - packed_accessor_kernel<<<1, 12>>>(foo_a, &trace); - -In addition to ``PackedTensorAccessor64`` and ``packed_accessor64`` there are -also the corresponding ``PackedTensorAccessor32`` and ``packed_accessor32`` -which use 32-bit integers for indexing. This can be quite a bit faster on CUDA -but may lead to overflows in the indexing calculations. - -Note that the template can hold other parameters such as the pointer restriction -and the integer type for indexing. See documentation for a thorough template -description of *accessors* and *packed accessors*. - - -Using Externally Created Data ------------------------------ - -If you already have your tensor data allocated in memory (CPU or CUDA), -you can view that memory as a ``Tensor`` in ATen: - -.. code-block:: cpp - - float data[] = { 1, 2, 3, - 4, 5, 6 }; - torch::Tensor f = torch::from_blob(data, {2, 3}); - -These tensors cannot be resized because ATen does not own the memory, but -otherwise behave as normal tensors. - -Scalars and zero-dimensional tensors ------------------------------------- - -In addition to the ``Tensor`` objects, ATen also includes ``Scalar``\s that -represent a single number. Like a Tensor, Scalars are dynamically typed and can -hold any one of ATen's number types. Scalars can be implicitly constructed from -C++ number types. Scalars are needed because some functions like ``addmm`` take -numbers along with Tensors and expect these numbers to be the same dynamic type -as the tensor. They are also used in the API to indicate places where a function -will *always* return a Scalar value, like ``sum``. - -.. code-block:: cpp - - namespace torch { - Tensor addmm(Scalar beta, const Tensor & self, - Scalar alpha, const Tensor & mat1, - const Tensor & mat2); - Scalar sum(const Tensor & self); - } // namespace torch - - // Usage. - torch::Tensor a = ... - torch::Tensor b = ... - torch::Tensor c = ... - torch::Tensor r = torch::addmm(1.0, a, .5, b, c); - -In addition to ``Scalar``\s, ATen also allows ``Tensor`` objects to be -zero-dimensional. These Tensors hold a single value and they can be references -to a single element in a larger ``Tensor``. They can be used anywhere a -``Tensor`` is expected. They are normally created by operators like `select` -which reduce the dimensions of a ``Tensor``. - -.. code-block:: cpp - - torch::Tensor two = torch::rand({10, 20}); - two[1][2] = 4; - // ^^^^^^ <- zero-dimensional Tensor diff --git a/docs/cpp/source/notes/tensor_creation.rst b/docs/cpp/source/notes/tensor_creation.rst deleted file mode 100644 index 25c9a8ca0cd46..0000000000000 --- a/docs/cpp/source/notes/tensor_creation.rst +++ /dev/null @@ -1,353 +0,0 @@ -Tensor Creation API -=================== - -This note describes how to create tensors in the PyTorch C++ API. It highlights -the available factory functions, which populate new tensors according to some -algorithm, and lists the options available to configure the shape, data type, -device and other properties of a new tensor. - -Factory Functions ------------------ - -A *factory function* is a function that produces a new tensor. There are many -factory functions available in PyTorch (both in Python and C++), which differ -in the way they initialize a new tensor before returning it. All factory -functions adhere to the following general "schema": - -.. code-block:: cpp - - torch::(, , ) - -Let's bisect the various parts of this "schema": - -1. ```` is the name of the function you would like to invoke, -2. ```` are any required or optional parameters a particular factory function accepts, -3. ```` is an object of type ``IntArrayRef`` and specifies the shape of the resulting tensor, -4. ```` is an instance of ``TensorOptions`` and configures the data type, device, layout and other properties of the resulting tensor. - -Picking a Factory Function -************************** - -The following factory functions are available at the time of this writing (the -hyperlinks lead to the corresponding Python functions, since they often have -more eloquent documentation -- the options are the same in C++): - -- `arange `_: Returns a tensor with a sequence of integers, -- `empty `_: Returns a tensor with uninitialized values, -- `eye `_: Returns an identity matrix, -- `full `_: Returns a tensor filled with a single value, -- `linspace `_: Returns a tensor with values linearly spaced in some interval, -- `logspace `_: Returns a tensor with values logarithmically spaced in some interval, -- `ones `_: Returns a tensor filled with all ones, -- `rand `_: Returns a tensor filled with values drawn from a uniform distribution on ``[0, 1)``. -- `randint `_: Returns a tensor with integers randomly drawn from an interval, -- `randn `_: Returns a tensor filled with values drawn from a unit normal distribution, -- `randperm `_: Returns a tensor filled with a random permutation of integers in some interval, -- `zeros `_: Returns a tensor filled with all zeros. - -Specifying a Size -***************** - -Functions that do not require specific arguments by nature of how they fill the -tensor can be invoked with just a size. For example, the following line creates -a vector with 5 components, initially all set to 1: - -.. code-block:: cpp - - torch::Tensor tensor = torch::ones(5); - - -What if we wanted to instead create a ``3 x 5`` matrix, or a ``2 x 3 x 4`` -tensor? In general, an ``IntArrayRef`` -- the type of the size parameter of factory -functions -- is constructed by specifying the size along each dimension in -curly braces. For example, ``{2, 3}`` for a tensor (in this case matrix) with -two rows and three columns, ``{3, 4, 5}`` for a three-dimensional tensor, and -``{2}`` for a one-dimensional tensor with two components. In the one -dimensional case, you can omit the curly braces and just pass the single -integer like we did above. Note that the squiggly braces are just one way of -constructing an ``IntArrayRef``. You can also pass an ``std::vector`` and -a few other types. Either way, this means we can construct a three-dimensional -tensor filled with values from a unit normal distribution by writing: - -.. code-block:: cpp - - torch::Tensor tensor = torch::randn({3, 4, 5}); - assert(tensor.sizes() == std::vector{3, 4, 5}); - -``tensor.sizes()`` returns an ``IntArrayRef`` which can be compared against an -``std::vector``, and we can see that it contains the sizes we passed -to the tensor. You can also write ``tensor.size(i)`` to access a single dimension, -which is equivalent to but preferred over ``tensor.sizes()[i]``. - -Passing Function-Specific Parameters -************************************ - -Neither ``ones`` nor ``randn`` accept any additional parameters to change their -behavior. One function which does require further configuration is ``randint``, -which takes an upper bound on the value for the integers it generates, as well -as an optional lower bound, which defaults to zero. Here we create a ``5 x 5`` -square matrix with integers between 0 and 10: - -.. code-block:: cpp - - torch::Tensor tensor = torch::randint(/*high=*/10, {5, 5}); - -And here we raise the lower bound to 3: - -.. code-block:: cpp - - torch::Tensor tensor = torch::randint(/*low=*/3, /*high=*/10, {5, 5}); - -The inline comments ``/*low=*/`` and ``/*high=*/`` are not required of course, -but aid readability just like keyword arguments in Python. - -.. tip:: - - The main take-away is that the size always follows the function specific - arguments. - -.. attention:: - - Sometimes a function does not need a size at all. For example, the size of - the tensor returned by ``arange`` is fully specified by its function-specific - arguments -- the lower and upper bound of a range of integers. In that case - the function does not take a ``size`` parameter. - -Configuring Properties of the Tensor -************************************ - -The previous section discussed function-specific arguments. Function-specific -arguments can only change the values with which tensors are filled, and -sometimes the size of the tensor. They never change things like the data type -(e.g. ``float32`` or ``int64``) of the tensor being created, or whether it -lives in CPU or GPU memory. The specification of these properties is left to -the very last argument to every factory function: a ``TensorOptions`` object, -discussed below. - -``TensorOptions`` is a class that encapsulates the construction axes of a -Tensor. With *construction axis* we mean a particular property of a Tensor that -can be configured before its construction (and sometimes changed afterwards). -These construction axes are: - -- The ``dtype`` (previously "scalar type"), which controls the data type of the - elements stored in the tensor, -- The ``layout``, which is either strided (dense) or sparse, -- The ``device``, which represents a compute device on which a tensor is stored (like a CPU or CUDA GPU), -- The ``requires_grad`` boolean to enable or disable gradient recording for a tensor, - -If you are used to PyTorch in Python, these axes will sound very familiar. The -allowed values for these axes at the moment are: - -- For ``dtype``: ``kUInt8``, ``kInt8``, ``kInt16``, ``kInt32``, ``kInt64``, ``kFloat32`` and ``kFloat64``, -- For ``layout``: ``kStrided`` and ``kSparse``, -- For ``device``: Either ``kCPU``, or ``kCUDA`` (which accepts an optional device index), -- For ``requires_grad``: either ``true`` or ``false``. - -.. tip:: - - There exist "Rust-style" shorthands for dtypes, like ``kF32`` instead of - ``kFloat32``. See `here - `_ - for the full list. - - -An instance of ``TensorOptions`` stores a concrete value for each of these -axes. Here is an example of creating a ``TensorOptions`` object that represents -a 64-bit float, strided tensor that requires a gradient, and lives on CUDA -device 1: - -.. code-block:: cpp - - auto options = - torch::TensorOptions() - .dtype(torch::kFloat32) - .layout(torch::kStrided) - .device(torch::kCUDA, 1) - .requires_grad(true); - - -Notice how we use the '"builder"-style methods of ``TensorOptions`` to -construct the object piece by piece. If we pass this object as the last -argument to a factory function, the newly created tensor will have these -properties: - -.. code-block:: cpp - - torch::Tensor tensor = torch::full({3, 4}, /*value=*/123, options); - - assert(tensor.dtype() == torch::kFloat32); - assert(tensor.layout() == torch::kStrided); - assert(tensor.device().type() == torch::kCUDA); // or device().is_cuda() - assert(tensor.device().index() == 1); - assert(tensor.requires_grad()); - -Now, you may be thinking: do I really need to specify each axis for every new -tensor I create? Fortunately, the answer is "no", as **every axis has a default -value**. These defaults are: - -- ``kFloat32`` for the dtype, -- ``kStrided`` for the layout, -- ``kCPU`` for the device, -- ``false`` for ``requires_grad``. - -What this means is that any axis you omit during the construction of a -``TensorOptions`` object will take on its default value. For example, this is -our previous ``TensorOptions`` object, but with the ``dtype`` and ``layout`` -defaulted: - -.. code-block:: cpp - - auto options = torch::TensorOptions().device(torch::kCUDA, 1).requires_grad(true); - -In fact, we can even omit all axes to get an entirely defaulted -``TensorOptions`` object: - -.. code-block:: cpp - - auto options = torch::TensorOptions(); // or `torch::TensorOptions options;` - -A nice consequence of this is that the ``TensorOptions`` object we just spoke -so much about can be entirely omitted from any tensor factory call: - -.. code-block:: cpp - - // A 32-bit float, strided, CPU tensor that does not require a gradient. - torch::Tensor tensor = torch::randn({3, 4}); - torch::Tensor range = torch::arange(5, 10); - -But the sugar gets sweeter: In the API presented here so far, you may have -noticed that the initial ``torch::TensorOptions()`` is quite a mouthful to -write. The good news is that for every construction axis (dtype, layout, device -and ``requires_grad``), there is one *free function* in the ``torch::`` -namespace which you can pass a value for that axis. Each function then returns -a ``TensorOptions`` object preconfigured with that axis, but allowing even -further modification via the builder-style methods shown above. For example, - -.. code-block:: cpp - - torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32)) - -is equivalent to - -.. code-block:: cpp - - torch::ones(10, torch::dtype(torch::kFloat32)) - -and further instead of - -.. code-block:: cpp - - torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32).layout(torch::kStrided)) - -we can just write - -.. code-block:: cpp - - torch::ones(10, torch::dtype(torch::kFloat32).layout(torch::kStrided)) - -which saves us quite a bit of typing. What this means is that in practice, you -should barely, if ever, have to write out ``torch::TensorOptions``. Instead use -the ``torch::dtype()``, ``torch::device()``, ``torch::layout()`` and -``torch::requires_grad()`` functions. - -A final bit of convenience is that ``TensorOptions`` is implicitly -constructible from individual values. This means that whenever a function has a -parameter of type ``TensorOptions``, like all factory functions do, we can -directly pass a value like ``torch::kFloat32`` or ``torch::kStrided`` in place -of the full object. Therefore, when there is only a single axis we would like -to change compared to its default value, we can pass only that value. As such, -what was - -.. code-block:: cpp - - torch::ones(10, torch::TensorOptions().dtype(torch::kFloat32)) - -became - -.. code-block:: cpp - - torch::ones(10, torch::dtype(torch::kFloat32)) - -and can finally be shortened to - -.. code-block:: cpp - - torch::ones(10, torch::kFloat32) - -Of course, it is not possible to modify further properties of the -``TensorOptions`` instance with this short syntax, but if all we needed was to -change one property, this is quite practical. - -In conclusion, we can now compare how ``TensorOptions`` defaults, together with -the abbreviated API for creating ``TensorOptions`` using free functions, allow -tensor creation in C++ with the same convenience as in Python. Compare this -call in Python: - -.. code-block:: python - - torch.randn(3, 4, dtype=torch.float32, device=torch.device('cuda', 1), requires_grad=True) - -with the equivalent call in C++: - -.. code-block:: cpp - - torch::randn({3, 4}, torch::dtype(torch::kFloat32).device(torch::kCUDA, 1).requires_grad(true)) - -Pretty close! - -Conversion ----------- - -Just as we can use ``TensorOptions`` to configure how new tensors should be -created, we can also use ``TensorOptions`` to convert a tensor from one set of -properties to a new set of properties. Such a conversion usually creates a new -tensor and does not occur in-place. For example, if we have a ``source_tensor`` -created with - -.. code-block:: cpp - - torch::Tensor source_tensor = torch::randn({2, 3}, torch::kInt64); - -we can convert it from ``int64`` to ``float32``: - -.. code-block:: cpp - - torch::Tensor float_tensor = source_tensor.to(torch::kFloat32); - -.. attention:: - - The result of the conversion, ``float_tensor``, is a new tensor pointing to - new memory, unrelated to the source ``source_tensor``. - -We can then move it from CPU memory to GPU memory: - -.. code-block:: cpp - - torch::Tensor gpu_tensor = float_tensor.to(torch::kCUDA); - -If you have multiple CUDA devices available, the above code will copy the -tensor to the *default* CUDA device, which you can configure with a -``torch::DeviceGuard``. If no ``DeviceGuard`` is in place, this will be GPU -1. If you would like to specify a different GPU index, you can pass it to -the ``Device`` constructor: - -.. code-block:: cpp - - torch::Tensor gpu_two_tensor = float_tensor.to(torch::Device(torch::kCUDA, 1)); - -In the case of CPU to GPU copy and reverse, we can also configure the memory -copy to be *asynchronous* by passing ``/*non_blocking=*/false`` as the last -argument to ``to()``: - -.. code-block:: cpp - - torch::Tensor async_cpu_tensor = gpu_tensor.to(torch::kCPU, /*non_blocking=*/true); - -Conclusion ----------- - -This note hopefully gave you a good understanding of how to create and convert -tensors in an idiomatic fashion using the PyTorch C++ API. If you have any -further questions or suggestions, please use our `forum -`_ or `GitHub issues -`_ to get in touch. diff --git a/docs/cpp/source/notes/tensor_cuda_stream.rst b/docs/cpp/source/notes/tensor_cuda_stream.rst deleted file mode 100644 index 4940317713635..0000000000000 --- a/docs/cpp/source/notes/tensor_cuda_stream.rst +++ /dev/null @@ -1,276 +0,0 @@ -Tensor CUDA Stream API -====================== - -A `CUDA Stream`_ is a linear sequence of execution that belongs to a specific CUDA device. -The PyTorch C++ API supports CUDA streams with the CUDAStream class and useful helper functions to make streaming operations easy. -You can find them in `CUDAStream.h`_. This note provides more details on how to use Pytorch C++ CUDA Stream APIs. - -.. _CUDA Stream: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#streams -.. _CUDAStream.h: https://pytorch.org/cppdocs/api/file_c10_cuda_CUDAStream.h.html#file-c10-cuda-cudastream-h -.. _CUDAStreamGuard.h: https://pytorch.org/cppdocs/api/structc10_1_1cuda_1_1_c_u_d_a_stream_guard.html - -Acquiring CUDA stream -********************* - -Pytorch's C++ API provides the following ways to acquire CUDA stream: - -1. Acquire a new stream from the CUDA stream pool, streams are preallocated from the pool and returned in a round-robin fashion. - -.. code-block:: cpp - - CUDAStream getStreamFromPool(const bool isHighPriority = false, DeviceIndex device = -1); - -.. tip:: - - You can request a stream from the high priority pool by setting isHighPriority to true, or a stream for a specific device - by setting device index (defaulting to the current CUDA stream's device index). - -2. Acquire the default CUDA stream for the passed CUDA device, or for the current device if no device index is passed. - -.. code-block:: cpp - - CUDAStream getDefaultCUDAStream(DeviceIndex device_index = -1); - -.. tip:: - - The default stream is where most computation occurs when you aren't explicitly using streams. - -3. Acquire the current CUDA stream, for the CUDA device with index ``device_index``, or for the current device if no device index is passed. - -.. code-block:: cpp - - CUDAStream getCurrentCUDAStream(DeviceIndex device_index = -1); - -.. tip:: - - The current CUDA stream will usually be the default CUDA stream for the device, but it may be different if someone - called ``setCurrentCUDAStream`` or used ``StreamGuard`` or ``CUDAStreamGuard``. - - - -Set CUDA stream -*************** - -Pytorch's C++ API provides the following ways to set CUDA stream: - -1. Set the current stream on the device of the passed in stream to be the passed in stream. - -.. code-block:: cpp - - void setCurrentCUDAStream(CUDAStream stream); - -.. attention:: - - This function may have nothing to do with the current device. It only changes the current stream on the stream's device. - We recommend using ``CUDAStreamGuard``, instead, since it switches to the stream's device and makes it the current stream on that device. - ``CUDAStreamGuard`` will also restore the current device and stream when it's destroyed - -2. Use ``CUDAStreamGuard`` to switch to a CUDA stream within a scope, it is defined in `CUDAStreamGuard.h`_ - -.. tip:: - - Use ``CUDAMultiStreamGuard`` if you need to set streams on multiple CUDA devices. - -CUDA Stream Usage Examples -************************** - -1. Acquiring and setting CUDA stream on the same device - -.. code-block:: cpp - - // This example shows how to acquire and set CUDA stream on the same device. - // `at::cuda::setCurrentCUDAStream` is used to set current CUDA stream - - // create a tensor on device 0 - torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(torch::kCUDA)); - // get a new CUDA stream from CUDA stream pool on device 0 - at::cuda::CUDAStream myStream = at::cuda::getStreamFromPool(); - // set current CUDA stream from default stream to `myStream` on device 0 - at::cuda::setCurrentCUDAStream(myStream); - // sum() on tensor0 uses `myStream` as current CUDA stream - tensor0.sum(); - - // get the default CUDA stream on device 0 - at::cuda::CUDAStream defaultStream = at::cuda::getDefaultCUDAStream(); - // set current CUDA stream back to default CUDA stream on device 0 - at::cuda::setCurrentCUDAStream(defaultStream); - // sum() on tensor0 uses `defaultStream` as current CUDA stream - tensor0.sum(); - -.. code-block:: cpp - - // This example is the same as previous example, but explicitly specify device - // index and use CUDA stream guard to set current CUDA stream - - // create a tensor on device 0 - torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(torch::kCUDA)); - // get a new stream from CUDA stream pool on device 0 - at::cuda::CUDAStream myStream = at::cuda::getStreamFromPool(false, 0); - // set the current CUDA stream to `myStream` within the scope using CUDA stream guard - { - at::cuda::CUDAStreamGuard guard(myStream); - // current CUDA stream is `myStream` from here till the end of bracket. - // sum() on tensor0 uses `myStream` as current CUDA stream - tensor0.sum(); - } - // current CUDA stream is reset to default CUDA stream after CUDA stream guard is destroyed - // sum() on tensor0 uses default CUDA stream on device 0 as current CUDA stream - tensor0.sum(); - -.. attention:: - - Above code is running on the same CUDA device. `setCurrentCUDAStream` will always set current CUDA stream on current device, - but note that `setCurrentCUDAStream` actually set current stream on the device of passed in CUDA stream. - - -2. Acquiring and setting CUDA streams on multiple devices. - -.. code-block:: cpp - - // This example shows how to acquire and set CUDA stream on two devices. - - // acquire new CUDA streams from CUDA stream pool on device 0 and device 1 - at::cuda::CUDAStream myStream0 = at::cuda::getStreamFromPool(false, 0); - at::cuda::CUDAStream myStream1 = at::cuda::getStreamFromPool(false, 1); - - // set current CUDA stream to `myStream0` on device 0 - at::cuda::setCurrentCUDAStream(myStream0); - // set current CUDA stream to `myStream1` on device 1 - at::cuda::setCurrentCUDAStream(myStream1); - - // create a tensor on device 0, no need to specify device index since - // current device index is 0 - torch::Tensor tensor0 = torch::ones({2, 2}, torch::device(at::kCUDA)); - // sum() on tensor0 use `myStream0` as current CUDA stream on device 0 - tensor0.sum(); - - // change the current device index to 1 by using CUDA device guard within a bracket scope - { - at::cuda::CUDAGuard device_guard{1}; - // create a tensor on device 1 - torch::Tensor tensor1 = torch::ones({2, 2}, torch::device(at::kCUDA)); - // sum() on tensor 1 uses `myStream1` as current CUDA stream on device 1 - tensor1.sum(); - } - - // current device is reset to device 0 after device_guard is destroyed - - // acquire a new CUDA stream on device 1 - at::cuda::CUDAStream myStream1_1 = at::cuda::getStreamFromPool(false, 1); - // create a new tensor on device 1 - torch::Tensor tensor1 = torch::ones({2, 2}, torch::device({torch::kCUDA, 1})); - - // change the current device index to 1 and current CUDA stream on device 1 - // to `myStream1_1` using CUDA stream guard within a scope - { - at::cuda::CUDAStreamGuard stream_guard(myStream1_1); - // sum() on tensor1 use `myStream1_1` as current CUDA stream on device 1 - tensor1.sum(); - } - - // current device is reset to device 0 and current CUDA stream on device 1 is - // reset to `myStream1` - - // sum() on tensor1 uses `myStream1` as current CUDA stream on device 1 - tensor1.sum(); - - -3. Working with CUDA multistream guard - -.. code-block:: cpp - - // This example shows how to use CUDA multistream guard to set - // two streams on two devices at the same time. - - // create two tensor, one on device 0, one on device 1 - torch::Tensor tensor0 = torch::ones({2, 2}, torch::device({torch::kCUDA, 0})); - torch::Tensor tensor1 = torch::ones({2, 2}, torch::device({torch::kCUDA, 1})); - - // acquire new CUDA streams from CUDA stream pool on device 0 and device 1 - at::cuda::CUDAStream myStream0 = at::cuda::getStreamFromPool(false, 0); - at::cuda::CUDAStream myStream1 = at::cuda::getStreamFromPool(false, 1); - - // set current CUDA stream on device 0 to `myStream0` and - // set current CUDA stream on device 1 to `myStream1` CUDA using multistream guard - { - at::cuda::CUDAMultiStreamGuard multi_guard({myStream0, myStream1}); - - // sum() on tensor0 uses `myStream0` as current CUDA stream on device 0 - tensor0.sum(); - // sum() on tensor1 uses `myStream1` as current CUDA stream on device 1 - tensor1.sum(); - } - - // current CUDA stream on device 0 is reset to default CUDA stream on device 0 - // current CUDA stream on device 1 is reset to default CUDA stream on device 1 - - // sum() on tensor0 uses default CUDA stream as current CUDA stream on device 0 - tensor0.sum(); - // sum() on tensor1 uses default CUDA stream as current CUDA stream on device 1 - tensor1.sum(); - -.. attention:: - ``CUDAMultiStreamGuard`` does not change current device index, it only changes the stream on - each passed in stream's device. Other than scope controlling, this guard is equivalent to - calling ``setCurrentCUDAStream`` on each passed in stream. - -4. A skeleton example for handling CUDA streams on multiple devices - -.. code-block:: cpp - - // This is a skeleton example that shows how to handle CUDA streams on multiple devices - // Suppose you want to do work on the non-default stream on two devices simultaneously, and we - // already have streams on both devices in two vectors. The following code shows three ways - // of acquiring and setting the streams. - - // Usage 0: acquire CUDA stream and set current CUDA stream with `setCurrentCUDAStream` - // Create a CUDA stream vector `streams0` on device 0 - std::vector streams0 = - {at::cuda::getDefaultCUDAStream(), at::cuda::getStreamFromPool()}; - // set current stream as `streams0[0]` on device 0 - at::cuda::setCurrentCUDAStream(streams0[0]); - - // create a CUDA stream vector `streams1` on device using CUDA device guard - std::vector streams1; - { - // device index is set to 1 within this scope - at::cuda::CUDAGuard device_guard(1); - streams1.push_back(at::cuda::getDefaultCUDAStream()); - streams1.push_back(at::cuda::getStreamFromPool()); - } - // device index is reset to 0 after device_guard is destroyed - - // set current stream as `streams1[0]` on device 1 - at::cuda::setCurrentCUDAStream(streams1[0]); - - - // Usage 1: use CUDA device guard to change the current device index only - { - at::cuda::CUDAGuard device_guard(1); - - // current device index is changed to 1 within scope - // current CUDA stream is still `streams1[0]` on device 1, no change - } - // current device index is reset to 0 after `device_guard` is destroyed - - - // Usage 2: use CUDA stream guard to change both current device index and current CUDA stream. - { - at::cuda::CUDAStreamGuard stream_guard(streams1[1]); - - // current device index and current CUDA stream are set to 1 and `streams1[1]` within scope - } - // current device index and current CUDA stream are reset to 0 and `streams0[0]` after - // stream_guard is destroyed - - - // Usage 3: use CUDA multi-stream guard to change multiple streams on multiple devices - { - // This is the same as calling `torch::cuda::setCurrentCUDAStream` on both streams - at::cuda::CUDAMultiStreamGuard multi_guard({streams0[1], streams1[1]}); - - // current device index is not change, still 0 - // current CUDA stream on device 0 and device 1 are set to `streams0[1]` and `streams1[1]` - } - // current CUDA stream on device 0 and device 1 are reset to `streams0[0]` and `streams1[0]` - // after `multi_guard` is destroyed. diff --git a/docs/cpp/source/notes/tensor_indexing.rst b/docs/cpp/source/notes/tensor_indexing.rst deleted file mode 100644 index 013ab7a737c75..0000000000000 --- a/docs/cpp/source/notes/tensor_indexing.rst +++ /dev/null @@ -1,99 +0,0 @@ -Tensor Indexing API -=================== - -Indexing a tensor in the PyTorch C++ API works very similar to the Python API. -All index types such as ``None`` / ``...`` / integer / boolean / slice / tensor -are available in the C++ API, making translation from Python indexing code to C++ -very simple. The main difference is that, instead of using the ``[]``-operator -similar to the Python API syntax, in the C++ API the indexing methods are: - -- ``torch::Tensor::index`` (`link `_) -- ``torch::Tensor::index_put_`` (`link `_) - -It's also important to note that index types such as ``None`` / ``Ellipsis`` / ``Slice`` -live in the ``torch::indexing`` namespace, and it's recommended to put ``using namespace torch::indexing`` -before any indexing code for convenient use of those index types. - -Here are some examples of translating Python indexing code to C++: - -Getter ------- - -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| Python | C++ (assuming ``using namespace torch::indexing``) | -+==========================================================+======================================================================================+ -| ``tensor[None]`` | ``tensor.index({None})`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[Ellipsis, ...]`` | ``tensor.index({Ellipsis, "..."})`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[1, 2]`` | ``tensor.index({1, 2})`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[True, False]`` | ``tensor.index({true, false})`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[1::2]`` | ``tensor.index({Slice(1, None, 2)})`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[torch.tensor([1, 2])]`` | ``tensor.index({torch::tensor({1, 2})})`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[..., 0, True, 1::2, torch.tensor([1, 2])]`` | ``tensor.index({"...", 0, true, Slice(1, None, 2), torch::tensor({1, 2})})`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ - -Setter ------- - -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| Python | C++ (assuming ``using namespace torch::indexing``) | -+==========================================================+======================================================================================+ -| ``tensor[None] = 1`` | ``tensor.index_put_({None}, 1)`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[Ellipsis, ...] = 1`` | ``tensor.index_put_({Ellipsis, "..."}, 1)`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[1, 2] = 1`` | ``tensor.index_put_({1, 2}, 1)`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[True, False] = 1`` | ``tensor.index_put_({true, false}, 1)`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[1::2] = 1`` | ``tensor.index_put_({Slice(1, None, 2)}, 1)`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[torch.tensor([1, 2])] = 1`` | ``tensor.index_put_({torch::tensor({1, 2})}, 1)`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ -| ``tensor[..., 0, True, 1::2, torch.tensor([1, 2])] = 1`` | ``tensor.index_put_({"...", 0, true, Slice(1, None, 2), torch::tensor({1, 2})}, 1)`` | -+----------------------------------------------------------+--------------------------------------------------------------------------------------+ - - -Translating between Python/C++ index types ------------------------------------------- - -The one-to-one translation between Python and C++ index types is as follows: - -+-------------------------+------------------------------------------------------------------------+ -| Python | C++ (assuming ``using namespace torch::indexing``) | -+=========================+========================================================================+ -| ``None`` | ``None`` | -+-------------------------+------------------------------------------------------------------------+ -| ``Ellipsis`` | ``Ellipsis`` | -+-------------------------+------------------------------------------------------------------------+ -| ``...`` | ``"..."`` | -+-------------------------+------------------------------------------------------------------------+ -| ``123`` | ``123`` | -+-------------------------+------------------------------------------------------------------------+ -| ``True`` | ``true`` | -+-------------------------+------------------------------------------------------------------------+ -| ``False`` | ``false`` | -+-------------------------+------------------------------------------------------------------------+ -| ``:`` or ``::`` | ``Slice()`` or ``Slice(None, None)`` or ``Slice(None, None, None)`` | -+-------------------------+------------------------------------------------------------------------+ -| ``1:`` or ``1::`` | ``Slice(1, None)`` or ``Slice(1, None, None)`` | -+-------------------------+------------------------------------------------------------------------+ -| ``:3`` or ``:3:`` | ``Slice(None, 3)`` or ``Slice(None, 3, None)`` | -+-------------------------+------------------------------------------------------------------------+ -| ``::2`` | ``Slice(None, None, 2)`` | -+-------------------------+------------------------------------------------------------------------+ -| ``1:3`` | ``Slice(1, 3)`` | -+-------------------------+------------------------------------------------------------------------+ -| ``1::2`` | ``Slice(1, None, 2)`` | -+-------------------------+------------------------------------------------------------------------+ -| ``:3:2`` | ``Slice(None, 3, 2)`` | -+-------------------------+------------------------------------------------------------------------+ -| ``1:3:2`` | ``Slice(1, 3, 2)`` | -+-------------------------+------------------------------------------------------------------------+ -| ``torch.tensor([1, 2])``| ``torch::tensor({1, 2})`` | -+-------------------------+------------------------------------------------------------------------+ diff --git a/docs/cpp/source/notes/versioning.rst b/docs/cpp/source/notes/versioning.rst deleted file mode 100644 index b7c79427bb031..0000000000000 --- a/docs/cpp/source/notes/versioning.rst +++ /dev/null @@ -1,29 +0,0 @@ -Library Versioning -================== - -We provide version number macros for identifying the version of LibTorch in use. -Example usage: - -.. code-block:: cpp - - #include - #include - - int main() { - std::cout << "PyTorch version from parts: " - << TORCH_VERSION_MAJOR << "." - << TORCH_VERSION_MINOR << "." - << TORCH_VERSION_PATCH << std::endl; - std::cout << "PyTorch version: " << TORCH_VERSION << std::endl; - } - -This will output something like: - -.. code-block:: text - - PyTorch version from parts: 1.8.0 - PyTorch version: 1.8.0 - -.. note:: - - These macros are only available in PyTorch >= 1.8.0. diff --git a/docs/cpp/source/stable.rst b/docs/cpp/source/stable.rst deleted file mode 100644 index c7c60995419da..0000000000000 --- a/docs/cpp/source/stable.rst +++ /dev/null @@ -1,388 +0,0 @@ -Torch Stable API -================ - -The PyTorch Stable C++ API provides a convenient high level interface to call -ABI-stable tensor operations and other utilities commonly used in custom operators. -These functions are designed to maintain binary compatibility across PyTorch versions, -making them suitable for use in ahead-of-time compiled code. - -For more information on the stable ABI, see the -`Stable ABI notes `_. - -Library Registration Macros ---------------------------- - -These macros provide stable ABI equivalents of the standard PyTorch operator -registration macros (``TORCH_LIBRARY``, ``TORCH_LIBRARY_IMPL``, etc.). -Use these when building custom operators that need to maintain binary -compatibility across PyTorch versions. - -``STABLE_TORCH_LIBRARY(ns, m)`` -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Defines a library of operators in a namespace using the stable ABI. - -This is the stable ABI equivalent of :c:macro:`TORCH_LIBRARY`. -Use this macro to define operator schemas that will maintain -binary compatibility across PyTorch versions. Only one ``STABLE_TORCH_LIBRARY`` -block can exist per namespace; use ``STABLE_TORCH_LIBRARY_FRAGMENT`` for -additional definitions in the same namespace from different translation units. - -**Parameters:** - -- ``ns`` - The namespace in which to define operators (e.g., ``mylib``). -- ``m`` - The name of the StableLibrary variable available in the block. - -**Example:** - -.. code-block:: cpp - - STABLE_TORCH_LIBRARY(mylib, m) { - m.def("my_op(Tensor input, int size) -> Tensor"); - m.def("another_op(Tensor a, Tensor b) -> Tensor"); - } - -Minimum compatible version: PyTorch 2.9. - -``STABLE_TORCH_LIBRARY_IMPL(ns, k, m)`` -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Registers operator implementations for a specific dispatch key using the stable ABI. - -This is the stable ABI equivalent of ``TORCH_LIBRARY_IMPL``. Use this macro -to provide implementations of operators for a specific dispatch key (e.g., -CPU, CUDA) while maintaining binary compatibility across PyTorch versions. - -.. note:: - - All kernel functions registered with this macro must be boxed using - the ``TORCH_BOX`` macro. - -**Parameters:** - -- ``ns`` - The namespace in which the operators are defined. -- ``k`` - The dispatch key (e.g., ``CPU``, ``CUDA``). -- ``m`` - The name of the StableLibrary variable available in the block. - -**Example:** - -.. code-block:: cpp - - STABLE_TORCH_LIBRARY_IMPL(mylib, CPU, m) { - m.impl("my_op", TORCH_BOX(&my_cpu_kernel)); - } - - STABLE_TORCH_LIBRARY_IMPL(mylib, CUDA, m) { - m.impl("my_op", TORCH_BOX(&my_cuda_kernel)); - } - -Minimum compatible version: PyTorch 2.9. - -``STABLE_TORCH_LIBRARY_FRAGMENT(ns, m)`` -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Extends operator definitions in an existing namespace using the stable ABI. - -This is the stable ABI equivalent of ``TORCH_LIBRARY_FRAGMENT``. Use this macro -to add additional operator definitions to a namespace that was already -created with ``STABLE_TORCH_LIBRARY``. - -**Parameters:** - -- ``ns`` - The namespace to extend. -- ``m`` - The name of the StableLibrary variable available in the block. - -Minimum compatible version: PyTorch 2.9. - -``TORCH_BOX(&func)`` -^^^^^^^^^^^^^^^^^^^ - -Wraps a function to conform to the stable boxed kernel calling convention. - -This macro takes an unboxed kernel function pointer and generates a boxed wrapper -that can be registered with the stable library API. - -**Parameters:** - -- ``func`` - The unboxed kernel function to wrap. - -**Example:** - -.. code-block:: cpp - - Tensor my_kernel(const Tensor& input, int64_t size) { - return input.reshape({size}); - } - - STABLE_TORCH_LIBRARY_IMPL(my_namespace, CPU, m) { - m.impl("my_op", TORCH_BOX(&my_kernel)); - } - -Minimum compatible version: PyTorch 2.9. - -Tensor Class ------------- - -The ``torch::stable::Tensor`` class offers a user-friendly C++ interface similar -to ``torch::Tensor`` while maintaining binary compatibility across PyTorch versions. - -.. doxygenclass:: torch::stable::Tensor - :members: - - -Device Class ------------- - -The ``torch::stable::Device`` class provides a user-friendly C++ interface similar -to ``c10::Device`` while maintaining binary compatibility across PyTorch versions. -It represents a compute device (CPU, CUDA, etc.) with an optional device index. - -.. doxygenclass:: torch::stable::Device - :members: - -DeviceGuard Class ------------------ - -The ``torch::stable::accelerator::DeviceGuard`` provides a user-friendly C++ -interface similar to ``c10::DeviceGuard`` while maintaining binary compatibility -across PyTorch versions. - -.. doxygenclass:: torch::stable::accelerator::DeviceGuard - :members: - -.. doxygenfunction:: torch::stable::accelerator::getCurrentDeviceIndex - - -Stream Utilities ----------------- - -For CUDA stream access, we currently recommend the ABI stable C shim API. This -will be improved in a future release with a more ergonomic wrapper. - -Getting the Current CUDA Stream -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -To obtain the current ``cudaStream_t`` for use in CUDA kernels: - -.. code-block:: cpp - - #include - #include - - // For now, we rely on the ABI stable C shim API to get the current CUDA stream. - // This will be improved in a future release. - // When using a C shim API, we need to use TORCH_ERROR_CODE_CHECK to - // check the error code and throw an appropriate runtime_error otherwise. - void* stream_ptr = nullptr; - TORCH_ERROR_CODE_CHECK( - aoti_torch_get_current_cuda_stream(tensor.get_device_index(), &stream_ptr)); - cudaStream_t stream = static_cast(stream_ptr); - - // Now you can use 'stream' in your CUDA kernel launches - my_kernel<<>>(args...); - -.. note:: - - The ``TORCH_ERROR_CODE_CHECK`` macro is required when using C shim APIs - to properly check error codes and throw appropriate exceptions. - -CUDA Error Checking Macros --------------------------- - -These macros provide stable ABI equivalents for CUDA error checking. -They wrap CUDA API calls and kernel launches, providing detailed error -messages using PyTorch's error formatting. - -``STD_CUDA_CHECK(EXPR)`` -^^^^^^^^^^^^^^^^^^^^^^^^ - -Checks the result of a CUDA API call and throws an exception on error. -Users of this macro are expected to include ``cuda_runtime.h``. - -**Example:** - -.. code-block:: cpp - - STD_CUDA_CHECK(cudaMalloc(&ptr, size)); - STD_CUDA_CHECK(cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost)); - -Minimum compatible version: PyTorch 2.10. - -``STD_CUDA_KERNEL_LAUNCH_CHECK()`` -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -Checks for errors from the most recent CUDA kernel launch. Equivalent to -``STD_CUDA_CHECK(cudaGetLastError())``. - -**Example:** - -.. code-block:: cpp - - my_kernel<<>>(args...); - STD_CUDA_KERNEL_LAUNCH_CHECK(); - -Minimum compatible version: PyTorch 2.10. - -Header-Only Utilities ---------------------- - -The ``torch::headeronly`` namespace provides header-only versions of common -PyTorch types and utilities. These can be used without linking against libtorch, -making them ideal for maintaining binary compatibility across PyTorch versions. - -Error Checking -^^^^^^^^^^^^^^ - -``STD_TORCH_CHECK`` is a header-only macro for runtime assertions: - -.. code-block:: cpp - - #include - - STD_TORCH_CHECK(condition, "Error message with ", variable, " interpolation"); - -Core Types -^^^^^^^^^^ - -The following ``c10::`` types are available as header-only versions under -``torch::headeronly::``: - -- ``torch::headeronly::ScalarType`` - Tensor data types (Float, Double, Int, etc.) -- ``torch::headeronly::DeviceType`` - Device types (CPU, CUDA, etc.) -- ``torch::headeronly::MemoryFormat`` - Memory layout formats (Contiguous, ChannelsLast, etc.) -- ``torch::headeronly::Layout`` - Tensor layouts (Strided, Sparse, etc.) - -.. code-block:: cpp - - #include - #include - #include - #include - - auto dtype = torch::headeronly::ScalarType::Float; - auto device_type = torch::headeronly::DeviceType::CUDA; - auto memory_format = torch::headeronly::MemoryFormat::Contiguous; - auto layout = torch::headeronly::Layout::Strided; - -TensorAccessor -^^^^^^^^^^^^^^ - -``TensorAccessor`` provides efficient, bounds-checked access to tensor data. -You can construct one from a stable tensor's data pointer, sizes, and strides: - -.. code-block:: cpp - - #include - - // Create a TensorAccessor for a 2D float tensor - auto sizes = tensor.sizes(); - auto strides = tensor.strides(); - torch::headeronly::TensorAccessor accessor( - static_cast(tensor.mutable_data_ptr()), - sizes.data(), - strides.data()); - - // Access elements - float value = accessor[i][j]; - -Dispatch Macros -^^^^^^^^^^^^^^^ - -Header-only dispatch macros (THO = Torch Header Only) are available for -dtype and device dispatching: - -.. code-block:: cpp - - #include - - THO_DISPATCH_FLOATING_TYPES(tensor.scalar_type(), "my_kernel", [&] { - // scalar_t is the resolved type - auto* data = tensor.data_ptr(); - }); - -Full API List -^^^^^^^^^^^^^ - -For the complete list of header-only APIs, see ``torch/header_only_apis.txt`` -in the PyTorch source tree. - -Stable Operators ----------------- - -Tensor Creation -^^^^^^^^^^^^^^^ - -.. doxygenfunction:: torch::stable::empty - -.. doxygenfunction:: torch::stable::empty_like - -.. doxygenfunction:: torch::stable::new_empty(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef size, std::optional dtype, std::optional layout, std::optional device, std::optional pin_memory) - -.. doxygenfunction:: torch::stable::new_zeros(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef size, std::optional dtype, std::optional layout, std::optional device, std::optional pin_memory) - -.. doxygenfunction:: torch::stable::full - -.. doxygenfunction:: torch::stable::from_blob - -Tensor Manipulation -^^^^^^^^^^^^^^^^^^^ - -.. doxygenfunction:: torch::stable::clone - -.. doxygenfunction:: torch::stable::contiguous - -.. doxygenfunction:: torch::stable::reshape - -.. doxygenfunction:: torch::stable::view - -.. doxygenfunction:: torch::stable::flatten - -.. doxygenfunction:: torch::stable::squeeze - -.. doxygenfunction:: torch::stable::unsqueeze - -.. doxygenfunction:: torch::stable::transpose - -.. doxygenfunction:: torch::stable::select - -.. doxygenfunction:: torch::stable::narrow - -.. doxygenfunction:: torch::stable::pad - - -Device and Type Conversion -^^^^^^^^^^^^^^^^^^^^^^^^^^ - -.. doxygenfunction:: torch::stable::to(const torch::stable::Tensor &self, std::optional dtype, std::optional layout, std::optional device, std::optional pin_memory, bool non_blocking, bool copy, std::optional memory_format) - -.. doxygenfunction:: torch::stable::to(const torch::stable::Tensor &self, torch::stable::Device device, bool non_blocking, bool copy) - -.. doxygenfunction:: torch::stable::fill_ - -.. doxygenfunction:: torch::stable::zero_ - -.. doxygenfunction:: torch::stable::copy_ - -.. doxygenfunction:: torch::stable::matmul - -.. doxygenfunction:: torch::stable::amax(const torch::stable::Tensor &self, int64_t dim, bool keepdim) - -.. doxygenfunction:: torch::stable::amax(const torch::stable::Tensor &self, torch::headeronly::IntHeaderOnlyArrayRef dims, bool keepdim) - -.. doxygenfunction:: torch::stable::sum - -.. doxygenfunction:: torch::stable::sum_out - -.. doxygenfunction:: torch::stable::subtract - -.. doxygenfunction:: torch::stable::parallel_for - -.. doxygenfunction:: torch::stable::get_num_threads - - -Parallelization Utilities -^^^^^^^^^^^^^^^^^^^^^^^^^ - -.. doxygenfunction:: torch::stable::parallel_for - -.. doxygenfunction:: torch::stable::get_num_threads diff --git a/docs/source/accelerator.md b/docs/source/accelerator.md index 53434cdfbfe25..1af297d315775 100644 --- a/docs/source/accelerator.md +++ b/docs/source/accelerator.md @@ -2,6 +2,7 @@ ```{eval-rst} .. automodule:: torch.accelerator + :no-members: ``` ```{eval-rst} @@ -20,6 +21,7 @@ set_device_idx current_device_index current_device_idx + get_device_capability set_stream current_stream synchronize @@ -50,6 +52,7 @@ :nosignatures: empty_cache + empty_host_cache get_memory_info max_memory_allocated max_memory_reserved diff --git a/docs/source/backends.md b/docs/source/backends.md index d85c53f200054..01faa514cdf48 100644 --- a/docs/source/backends.md +++ b/docs/source/backends.md @@ -24,6 +24,7 @@ These backends include: - `torch.backends.nnpack` - `torch.backends.openmp` - `torch.backends.opt_einsum` +- `torch.backends.python_native` - `torch.backends.xeon` ## torch.backends.cpu @@ -105,6 +106,10 @@ These backends include: .. autofunction:: torch.backends.cuda.preferred_rocm_fa_library ``` +```{eval-rst} +.. autofunction:: torch.backends.cuda.is_ck_sdpa_available +``` + ```{eval-rst} .. autofunction:: torch.backends.cuda.preferred_linalg_library ``` @@ -396,6 +401,122 @@ These backends include: ``` +## torch.backends.python_native + +```{eval-rst} +.. automodule:: torch.backends.python_native +``` + +The `torch.backends.python_native` module provides user control over native operators implemented in python +via. DSLs (Domain Specific Languages) that are defined in `torch._native`. This allows users to selectively +enable or disable high-performance implementations from various DSLs like Triton and CuteDSL. + +### Module-level Functions + +```{eval-rst} +.. autofunction:: torch.backends.python_native.get_dsl_operations +``` + +```{eval-rst} +.. autofunction:: torch.backends.python_native.disable_operations +``` + +```{eval-rst} +.. autofunction:: torch.backends.python_native.enable_operations +``` + +```{eval-rst} +.. autofunction:: torch.backends.python_native.disable_dispatch_keys +``` + +```{eval-rst} +.. autofunction:: torch.backends.python_native.enable_dispatch_keys +``` + +```{eval-rst} +.. autofunction:: torch.backends.python_native.operations_disabled +``` + +### Module-level Properties + +```{eval-rst} +.. attribute:: available_dsls + + A :class:`list` of :class:`str` containing the names of DSLs that are available at runtime. + This is a subset of :attr:`all_dsls` that have their runtime dependencies satisfied. +``` + +```{eval-rst} +.. attribute:: all_dsls + + A :class:`list` of :class:`str` containing the names of all registered DSLs, whether + available at runtime or not. +``` + +### DSL Controllers + +For each registered DSL (e.g., `triton`, `cutedsl`), auto-populated controller modules are available: + +```{eval-rst} +.. currentmodule:: torch.backends.python_native +``` + +#### DSL Properties + +Each DSL controller (e.g., `torch.backends.python_native.triton`) provides the following properties: + +| Property | Type | Description | +|----------|------|-------------| +| `name` | `str` | The name of the DSL | +| `available` | `bool` | Whether the DSL's runtime dependencies are available | +| `enabled` | `bool` | Controls whether all operations from this DSL are enabled. Setting to `False` disables all operations from the DSL, while `True` re-enables them | +| `version` | `Version` or `None` | The version of the DSL runtime, if available. Returns `None` if the DSL is not available | + +#### DSL Methods + +Each DSL controller provides the following methods: + +**disable()** + Disable all operations from this DSL. + +**enable()** + Re-enable all operations from this DSL. + +**disabled()** + Context manager that temporarily disables all operations from this DSL. + Operations are automatically re-enabled when exiting the context. + + Example:: + + with torch.backends.python_native.triton.disabled(): + # Triton operations are disabled here + result = model(input) + # Triton operations restored here + +### Usage Examples + +```{eval-rst} +.. code-block:: python + + import torch.backends.python_native as pn + + # Query available DSLs + print(pn.available_dsls) # ['triton', 'cutedsl'] + + # Disable all Triton operations + pn.triton.enabled = False + + # Temporarily disable CuteDSL operations + with pn.cutedsl.disabled(): + result = model(input) # CuteDSL ops disabled + + # Disable specific operations across all DSLs + pn.disable_operations('scaled_mm', '_flash_attention_forward') + + # Query operations for a specific DSL + triton_ops = pn.get_dsl_operations('triton') +``` + ## torch.backends.xeon ```{eval-rst} diff --git a/docs/source/community/persons_of_interest.rst b/docs/source/community/persons_of_interest.rst index 4b6c553a5dcd1..086a3e4c87619 100644 --- a/docs/source/community/persons_of_interest.rst +++ b/docs/source/community/persons_of_interest.rst @@ -346,6 +346,7 @@ ExecuTorch (Edge, Mobile) - Mergen Nachin (`mergennachin `__) - Jacob Szwejbka (`JacobSzwejbka `__) +- Digant Desai (`digantdesai `__) - (emeritus) Kimish Patel (`kimishpatel `__) - (emeritus) Dave Bort (`dbort `__) - (emeritus) Martin Yuan (`iseeyuan `__) diff --git a/docs/source/conf.py b/docs/source/conf.py index eae6ac07dfea2..9203ceefd8a84 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -78,6 +78,12 @@ "html_image", ] +# Don't execute notebooks during the docs build. Notebook correctness is +# verified by the separate docs_test CI job; re-executing them here just +# adds ~3 minutes to the build for no benefit. +nb_execution_mode = "off" +suppress_warnings = ["myst-nb.lexer"] + html_baseurl = "https://docs.pytorch.org/docs/stable/" # needed for sphinx-sitemap sitemap_locales = [None] sitemap_excludes = [ @@ -253,31 +259,18 @@ coverage_ignore_functions = [ "main", # utility script "run", # utility script - # torch.distributions.constraints - "is_dependent", # torch.hub "import_module", - # torch.jit - "export_opnames", # torch.jit.unsupported_tensor_ops "execWrapper", - # torch.onnx - "unregister_custom_op_symbolic", - # torch.ao.quantization - "default_eval_fn", # torch.backends "disable_global_flags", "flags_frozen", - # torch.distributed.algorithms.ddp_comm_hooks - "register_ddp_comm_hook", # torch.nn.parallel "DistributedDataParallelCPU", # torch.utils - "set_module", "burn_in_info", "get_info_and_burn_skeleton", - "get_inline_skeleton", - "get_model_info", "get_storage_info", "hierarchical_pickle", # torch.amp.autocast_mode @@ -612,70 +605,33 @@ "make_sharded_output_tensor", # torch.fx.passes.annotate_getitem_nodes "annotate_getitem_nodes", - # torch.fx.passes.backends.cudagraphs - "partition_cudagraphs", # torch.fx.passes.dialect.common.cse_pass "get_CSE_banned_ops", # torch.fx.passes.graph_manipulation "get_size_of_all_nodes", "get_size_of_node", "get_tensor_meta", - "replace_target_nodes_with", - # torch.fx.passes.infra.pass_manager - "pass_result_wrapper", - "this_before_that_pass_constraint", - # torch.fx.passes.operator_support - "any_chain", - "chain", - "create_op_support", - # torch.fx.passes.param_fetch - "default_matching", - "extract_attrs_for_lowering", - "lift_lowering_attrs_to_nodes", - # torch.fx.passes.pass_manager - "inplace_wrapper", - "log_hook", - "loop_pass", - "these_before_those_pass_constraint", - "this_before_that_pass_constraint", - # torch.fx.passes.regional_inductor - "regional_inductor", - # torch.fx.passes.reinplace - "reinplace", # torch.fx.passes.split_module "split_module", # torch.fx.passes.split_utils "getattr_recursive", - "setattr_recursive", "split_by_tags", # torch.fx.passes.splitter_base "generate_inputs_for_submodules", # torch.fx.passes.tools_common "get_acc_ops_name", "get_node_target", - "is_node_output_tensor", "legalize_graph", # torch.fx.passes.utils.common - "compare_graphs", "lift_subgraph_as_module", # torch.fx.passes.utils.fuser_utils - "erase_nodes", "fuse_as_graphmodule", "fuse_by_partitions", "insert_subgm", - "topo_sort", - "validate_partition", # torch.fx.passes.utils.source_matcher_utils - "check_subgraphs_connected", "get_source_partitions", # torch.fx.proxy "assert_fn", - # torch.fx.subgraph_rewriter - "replace_pattern", - "replace_pattern_with_filters", - # torch.fx.tensor_type - "is_consistent", - "is_more_precise", # torch.fx.traceback "format_stack", "get_current_meta", @@ -683,9 +639,9 @@ "preserve_node_meta", "reset_grad_fn_seq_nr", "set_current_meta", + "set_current_replay_node", "set_grad_fn_seq_nr", "set_stack_trace", - "set_current_replay_node", "get_current_replay_node", # torch.jit.annotations "ann_to_type", @@ -726,13 +682,6 @@ "quantize_linear_modules", "quantize_rnn_cell_modules", "quantize_rnn_modules", - # torch.library - "define", - "get_ctx", - "impl", - "impl_abstract", - # torch.masked.maskedtensor.core - "is_masked_tensor", # torch.masked.maskedtensor.creation "as_masked_tensor", "masked_tensor", @@ -1018,19 +967,8 @@ "check_export_model_diff", "verify", "verify_aten_graph", - # torch.optim.optimizer - "register_optimizer_step_post_hook", - "register_optimizer_step_pre_hook", # torch.overrides "enable_reentrant_dispatch", - # torch.package.analyze.find_first_use_of_broken_modules - "find_first_use_of_broken_modules", - # torch.package.analyze.is_from_package - "is_from_package", - # torch.package.analyze.trace_dependencies - "trace_dependencies", - # torch.profiler.itt - "range", # torch.profiler.profiler "schedule", "supported_activities", @@ -1038,8 +976,6 @@ # torch.return_types "pytree_register_structseq", # torch.serialization - "check_module_version_greater_or_equal", - "default_restore_location", "load", "location_tag", "mkdtemp", @@ -1060,8 +996,6 @@ "hann", "kaiser", "nuttall", - # torch.sparse.semi_structured - "to_sparse_semi_structured", # torch.utils.backend_registration "generate_methods_for_privateuse1_backend", "rename_privateuse1_backend", @@ -2502,9 +2436,123 @@ def setup(app): app.connect("autodoc-process-docstring", process_docstring) app.connect("html-page-context", hide_edit_button_for_pages) app.config.add_last_updated = True + + # Force serial reads to avoid pipe congestion from large env pickles. + # Sphinx's parallel read sends the entire environment (100s of MB for + # PyTorch) through a 64KB OS pipe per worker, which causes extreme + # slowdowns with many workers. Serial reads avoid this overhead while + # parallel writes (which send trivial payloads) remain enabled. + from sphinx.builders import Builder + + _orig_read_serial = Builder._read_serial + + def _serial_read_ignoring_nproc(self, docnames, nproc=1): + return _orig_read_serial(self, docnames) + + Builder._read_parallel = _serial_read_ignoring_nproc + + # Skip pickling doctrees to disk for the HTML builder. This is only used + # for incremental rebuilds which don't apply in CI clean builds. Saves + # ~2 minutes by avoiding serializing large autodoc-generated doctrees. + # Other builders (doctest, coverage) may need doctrees on disk. + from sphinx.builders.html import StandaloneHTMLBuilder + + def _write_doctree_no_disk(self, docname, doctree, *, _cache=True): + # Still do the cleanup and in-memory caching, just skip the disk I/O + doctree.reporter = None + doctree.transformer = None + doctree.settings = doctree.settings.copy() + doctree.settings.warning_stream = None + doctree.settings.env = None + from docutils.utils import DependencyList + + doctree.settings.record_dependencies = DependencyList() + if _cache: + self.env._write_doc_doctree_cache[docname] = doctree + + StandaloneHTMLBuilder.write_doctree = _write_doctree_no_disk + + _skip_git_dates_on_ci(app) + _fix_katex_server_race(app) + return {"version": "0.1", "parallel_read_safe": True} +def _fix_katex_server_race(app): + """Retry on ConnectionRefusedError when connecting to the KaTeX server. + + sphinxcontrib-katex 0.9.x starts a Node.js server and connects via Unix + socket. There's a race: Python sees the socket file (created by bind()) + and calls connect() before Node.js has called listen(). On slow CI + machines this causes ConnectionRefusedError. Fix by retrying connect(). + """ + try: + from sphinxcontrib.katex import KaTeXServer + except ImportError: + return + + import socket as _socket + import time as _time + + @classmethod + def _start_with_retry(cls, rundir, timeout): + from subprocess import PIPE, Popen + + from sphinxcontrib.katex import ONE_MILLISECOND + + socket_path = rundir / "katex.sock" + cmd = cls.build_command(socket=socket_path) + process = Popen(cmd, stdin=PIPE, stdout=PIPE, cwd=rundir) + + startup_start = _time.monotonic() + while not socket_path.is_socket(): + _time.sleep(ONE_MILLISECOND) + if _time.monotonic() - startup_start > timeout: + raise cls.timeout_error(timeout) + + # Retry connect() — bind() creates the socket file but listen() + # is async in Node.js and may not be ready yet. + while True: + remaining = startup_start + timeout - _time.monotonic() + if remaining <= 0: + raise cls.timeout_error(timeout) + try: + sock = _socket.socket(_socket.AF_UNIX, _socket.SOCK_STREAM) + sock.settimeout(remaining) + sock.connect(str(socket_path)) + break + except ConnectionRefusedError: + sock.close() + _time.sleep(ONE_MILLISECOND) + except TimeoutError: + sock.close() + raise cls.timeout_error(timeout) # noqa: B904 + + return process, sock + + KaTeXServer.start_server_process = _start_with_retry + + +def _skip_git_dates_on_ci(app): + """Skip git date lookups on non-release CI builds. + + The pytorch theme runs 2 git subprocess calls per page to display + "Created/Updated" dates. On CI preview builds this is wasteful (dates + aren't needed) and problematic (treeless clones make git log slow). + Release builds (WITH_PUSH=true) keep the original behavior so dates + appear in published docs. + """ + if not os.environ.get("CI") or os.environ.get("WITH_PUSH") == "true": + return + + try: + import pytorch_sphinx_theme2 + except ImportError: + return + + pytorch_sphinx_theme2.get_git_dates = lambda path: ("Unknown", "Unknown") + + def hide_edit_button_for_pages(app, pagename, templatename, context, doctree): if pagename.startswith("generated/"): context["theme_use_edit_page_button"] = False diff --git a/docs/source/cuda.md b/docs/source/cuda.md index fc30c4f851d92..929f90a823bc7 100644 --- a/docs/source/cuda.md +++ b/docs/source/cuda.md @@ -165,6 +165,7 @@ :toctree: generated :nosignatures: + caching_allocator_disabled caching_allocator_enable ``` @@ -307,7 +308,19 @@ See the docs for {class}`~torch.cuda.green_contexts.GreenContext` for an example ``` ```{eval-rst} -.. py:module:: torch.cuda.nvtx +.. automodule:: torch.cuda.nvtx +``` + +```{eval-rst} +.. currentmodule:: torch.cuda.nvtx +``` + +```{eval-rst} +.. autofunction:: range_start +``` + +```{eval-rst} +.. autofunction:: range_end ``` ```{eval-rst} diff --git a/docs/source/ddp_comm_hooks.md b/docs/source/ddp_comm_hooks.md index 059c388cd003a..5ead08d6d8564 100644 --- a/docs/source/ddp_comm_hooks.md +++ b/docs/source/ddp_comm_hooks.md @@ -33,6 +33,13 @@ Particularly, {class}`torch.distributed.GradBucket` represents a bucket of gradi .. autofunction:: torch.distributed.GradBucket.parameters ``` +## Register Communication Hook + +```{eval-rst} +.. currentmodule:: torch.distributed.algorithms.ddp_comm_hooks +.. autofunction:: register_ddp_comm_hook +``` + ## Default Communication Hooks Default communication hooks are simple **stateless** hooks, so the input state diff --git a/docs/source/distributed.md b/docs/source/distributed.md index b2699b8530624..6be381685f783 100644 --- a/docs/source/distributed.md +++ b/docs/source/distributed.md @@ -663,6 +663,10 @@ with torch.profiler(): Please refer to the [profiler documentation](https://pytorch.org/docs/main/profiler.html) for a full overview of profiler features. +```{eval-rst} +.. autofunction:: torch.distributed.distributed_c10d.record_comm +``` + ## Optimization with Symmetric Memory ### Copy Engine Collectives diff --git a/docs/source/distributed.tensor.md b/docs/source/distributed.tensor.md index 7361c04924675..d768de0e096eb 100644 --- a/docs/source/distributed.tensor.md +++ b/docs/source/distributed.tensor.md @@ -264,7 +264,10 @@ these features. ```{eval-rst} .. autofunction:: register_sharding +``` +```{eval-rst} +.. autofunction:: implicit_replication ``` % modules that are missing docs, add the doc later when necessary diff --git a/docs/source/fx.experimental.md b/docs/source/fx.experimental.md index e020adac03f8c..d842fa6a4edf3 100644 --- a/docs/source/fx.experimental.md +++ b/docs/source/fx.experimental.md @@ -362,6 +362,7 @@ These APIs are experimental and subject to change without notice. view_inference_rule register_inference_rule transpose_inference_rule + range_check ``` ## torch.fx.experimental.migrate_gradual_types.constraint_transformation diff --git a/docs/source/fx.md b/docs/source/fx.md index 426f19a5f6689..50181b3c189e8 100644 --- a/docs/source/fx.md +++ b/docs/source/fx.md @@ -8,6 +8,7 @@ ## Overview ```{eval-rst} .. automodule:: torch.fx + :no-members: ``` @@ -1176,6 +1177,236 @@ The set of leaf modules can be customized by overriding annotate_fn ``` +## torch.fx.subgraph_rewriter + +```{eval-rst} +.. currentmodule:: torch.fx.subgraph_rewriter +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + replace_pattern + replace_pattern_with_filters +``` + +## torch.fx.tensor_type + +```{eval-rst} +.. currentmodule:: torch.fx.tensor_type +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + is_consistent + is_more_precise +``` + +## torch.fx.passes.backends.cudagraphs + +```{eval-rst} +.. currentmodule:: torch.fx.passes.backends.cudagraphs +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + partition_cudagraphs +``` + +## torch.fx.passes.graph_manipulation + +```{eval-rst} +.. currentmodule:: torch.fx.passes.graph_manipulation +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + replace_target_nodes_with +``` + +## torch.fx.passes.infra.pass_manager + +```{eval-rst} +.. currentmodule:: torch.fx.passes.infra.pass_manager +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + pass_result_wrapper + this_before_that_pass_constraint +``` + +## torch.fx.passes.operator_support + +```{eval-rst} +.. currentmodule:: torch.fx.passes.operator_support +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + any_chain + chain + create_op_support +``` + +## torch.fx.passes.param_fetch + +```{eval-rst} +.. automodule:: torch.fx.passes.param_fetch + :no-members: +``` + +```{eval-rst} +.. currentmodule:: torch.fx.passes.param_fetch +``` + +```{eval-rst} +.. autofunction:: default_matching +``` + +```{eval-rst} +.. autofunction:: extract_attrs_for_lowering +``` + +```{eval-rst} +.. autofunction:: lift_lowering_attrs_to_nodes +``` + +## torch.fx.passes.pass_manager + +```{eval-rst} +.. currentmodule:: torch.fx.passes.pass_manager +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + inplace_wrapper + log_hook + loop_pass + these_before_those_pass_constraint + this_before_that_pass_constraint +``` + +## torch.fx.passes.regional_inductor + +```{eval-rst} +.. currentmodule:: torch.fx.passes.regional_inductor +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + regional_inductor +``` + +## torch.fx.passes.reinplace + +```{eval-rst} +.. currentmodule:: torch.fx.passes.reinplace +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + reinplace +``` + +## torch.fx.passes.split_utils + +```{eval-rst} +.. currentmodule:: torch.fx.passes.split_utils +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + setattr_recursive +``` + +## torch.fx.passes.tools_common + +```{eval-rst} +.. currentmodule:: torch.fx.passes.tools_common +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + is_node_output_tensor +``` + +## torch.fx.passes.utils.common + +```{eval-rst} +.. currentmodule:: torch.fx.passes.utils.common +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + compare_graphs +``` + +## torch.fx.passes.utils.fuser_utils + +```{eval-rst} +.. currentmodule:: torch.fx.passes.utils.fuser_utils +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + erase_nodes + topo_sort + validate_partition +``` + +## torch.fx.passes.utils.source_matcher_utils + +```{eval-rst} +.. currentmodule:: torch.fx.passes.utils.source_matcher_utils +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + check_subgraphs_connected +``` + @@ -1218,7 +1449,7 @@ The set of leaf modules can be customized by overriding .. py:module:: torch.fx.passes.infra.pass_manager .. py:module:: torch.fx.passes.net_min_base .. py:module:: torch.fx.passes.operator_support -.. py:module:: torch.fx.passes.param_fetch + .. py:module:: torch.fx.passes.pass_manager .. py:module:: torch.fx.passes.regional_inductor .. py:module:: torch.fx.passes.reinplace diff --git a/docs/source/masked.md b/docs/source/masked.md index 6a79a6be2ea95..af9bc9911720f 100644 --- a/docs/source/masked.md +++ b/docs/source/masked.md @@ -304,6 +304,20 @@ The following ops are currently supported: Tensor.view ``` +## torch.masked.maskedtensor.core + +```{eval-rst} +.. currentmodule:: torch.masked.maskedtensor.core +``` + +```{eval-rst} +.. autosummary:: + :toctree: generated + :nosignatures: + + is_masked_tensor +``` + ```{eval-rst} .. This module needs to be documented. Adding here in the meantime .. for tracking purposes diff --git a/docs/source/multiprocessing.md b/docs/source/multiprocessing.md index 6669fcaa24b30..74944ca58ed20 100644 --- a/docs/source/multiprocessing.md +++ b/docs/source/multiprocessing.md @@ -206,6 +206,12 @@ terminate processes upon detecting an error in one of them. % for tracking purposes +## torch.multiprocessing.pool + +```{eval-rst} +.. currentmodule:: torch.multiprocessing.pool +``` + ```{eval-rst} .. py:module:: torch.multiprocessing.pool ``` diff --git a/docs/source/notes/cuda.rst b/docs/source/notes/cuda.rst index 15c67f4ab67d8..399bf6a19bdb9 100644 --- a/docs/source/notes/cuda.rst +++ b/docs/source/notes/cuda.rst @@ -627,6 +627,24 @@ Available options: on all the CUDA devices to a specified fraction of the available memory. This is a value between 0 and 1. Attempting to allocate more memory will raise an out of memory error. +* ``throw_on_cudamalloc_oom`` (default: ``False``) + When set to ``True``, the allocator will preemptively reject allocations that would + exceed the ``per_process_memory_fraction`` limit by throwing an ``OutOfMemoryError`` + instead of attempting the allocation through the driver, which could trigger a fatal + GPU runtime abort (e.g., ``HSA_STATUS_ERROR_OUT_OF_RESOURCES`` on ROCm). The rejected + allocation skips the retry chain and throws immediately, allowing the caller to catch + the exception and handle it gracefully. + + Example configuration: + + .. code-block:: bash + + PYTORCH_CUDA_ALLOC_CONF=per_process_memory_fraction:0.95,throw_on_cudamalloc_oom:True + + This is particularly useful for inference serving, where a fatal GPU OOM would crash the + server process. With this option, the serving framework can catch the ``OutOfMemoryError`` + and reject the individual request while continuing to serve subsequent requests. + .. note:: Some stats reported by the diff --git a/docs/source/notes/ddp.rst b/docs/source/notes/ddp.rst index 07804f8bae166..4f74114fabce8 100644 --- a/docs/source/notes/ddp.rst +++ b/docs/source/notes/ddp.rst @@ -199,7 +199,7 @@ DistributedDataParallel parameters when ``find_unused_parameters`` is set to ``True`` in DDP constructor. -.. image:: https://user-images.githubusercontent.com/16999635/72313120-4e7c1c80-3658-11ea-9c6d-44336b2daeac.png +.. image:: https://github.com/user-attachments/assets/b6fe5258-f1f1-4c73-8e1d-6fd96406faa2 :alt: ddp_code.png :width: 400 px diff --git a/docs/source/optim.md b/docs/source/optim.md index c389fa6c37d2b..ffbc172ede073 100644 --- a/docs/source/optim.md +++ b/docs/source/optim.md @@ -149,6 +149,18 @@ for input, target in dataset: Optimizer.zero_grad ``` +## Module-level hooks + +```{eval-rst} +.. currentmodule:: torch.optim.optimizer + +.. autofunction:: register_optimizer_step_post_hook + +.. autofunction:: register_optimizer_step_pre_hook + +.. currentmodule:: torch.optim +``` + ## Algorithms ```{eval-rst} diff --git a/docs/source/package.md b/docs/source/package.md index 1b50f743d5793..dcc699cf545f6 100644 --- a/docs/source/package.md +++ b/docs/source/package.md @@ -726,12 +726,23 @@ statements more clearly show whether they are referring to packaged code or not. :members: ``` - +## Analysis Utilities ```{eval-rst} .. py:module:: torch.package.analyze.find_first_use_of_broken_modules .. py:module:: torch.package.analyze.is_from_package .. py:module:: torch.package.analyze.trace_dependencies + +.. currentmodule:: torch.package.analyze.find_first_use_of_broken_modules + +.. autofunction:: find_first_use_of_broken_modules + +.. currentmodule:: torch.package.analyze.is_from_package + +.. autofunction:: is_from_package + +.. currentmodule:: torch.package.analyze.trace_dependencies + +.. autofunction:: trace_dependencies .. py:module:: torch.package.file_structure_representation .. py:module:: torch.package.find_file_dependencies .. py:module:: torch.package.glob_group diff --git a/docs/source/profiler.md b/docs/source/profiler.md index 1578b7334d849..45789b95dea71 100644 --- a/docs/source/profiler.md +++ b/docs/source/profiler.md @@ -26,6 +26,8 @@ .. autofunction:: torch.profiler.schedule .. autofunction:: torch.profiler.tensorboard_trace_handler + +.. autofunction:: torch.profiler.supported_activities ``` ## Intel Instrumentation and Tracing Technology APIs @@ -38,6 +40,8 @@ .. autofunction:: torch.profiler.itt.range_push .. autofunction:: torch.profiler.itt.range_pop + +.. autofunction:: torch.profiler.itt.range ``` AOTInductorArrayRefTensor (zero-copy) +// ------------------------------------------------------------------------- +template +inline void arrayref_tensor_to_c( + const ArrayRefTensor& src, + AOTInductorArrayRefTensor& dst) { + const auto sizes = src.sizes(); + const auto strides = src.strides(); + dst.data = const_cast(static_cast(src.data())); + dst.numel = static_cast(src.numel()); + dst.ndim = static_cast(sizes.size()); + dst.dtype = aoti_torch_dtype>(); + dst.device_type = src.device_type(); + dst.device_idx = src.device_idx(); + + validate_arrayref_tensor_ndim(dst.ndim); + assert(dst.ndim <= AOTI_ARRAYREF_TENSOR_MAX_DIMS); + std::memcpy(dst.sizes, sizes.data(), dst.ndim * sizeof(int64_t)); + std::memcpy(dst.strides, strides.data(), dst.ndim * sizeof(int64_t)); + const int32_t remaining = AOTI_ARRAYREF_TENSOR_MAX_DIMS - dst.ndim; + std::memset(dst.sizes + dst.ndim, 0, remaining * sizeof(int64_t)); + std::memset(dst.strides + dst.ndim, 0, remaining * sizeof(int64_t)); + std::memset(dst.reserved, 0, sizeof(dst.reserved)); +} + +template +inline AOTInductorArrayRefTensor arrayref_tensor_to_c( + const ArrayRefTensor& src) { + AOTInductorArrayRefTensor dst; + arrayref_tensor_to_c(src, dst); + return dst; +} + +// ------------------------------------------------------------------------- +// AOTInductorArrayRefTensor --> ArrayRefTensor (zero-copy) +// ------------------------------------------------------------------------- +template +inline ArrayRefTensor c_to_arrayref_tensor( + const AOTInductorArrayRefTensor& src) { + validate_arrayref_tensor_ndim(src.ndim); + return ArrayRefTensor( + MiniArrayRef( + static_cast(const_cast(src.data)), + static_cast(src.numel)), + MiniArrayRef(src.sizes, static_cast(src.ndim)), + MiniArrayRef(src.strides, static_cast(src.ndim)), + src.device_type, + src.device_idx); +} + +} // namespace torch::aot_inductor diff --git a/torch/csrc/inductor/aoti_runtime/interface.h b/torch/csrc/inductor/aoti_runtime/interface.h index ffccdd94e5be2..a47bb3e4c67ff 100644 --- a/torch/csrc/inductor/aoti_runtime/interface.h +++ b/torch/csrc/inductor/aoti_runtime/interface.h @@ -35,6 +35,52 @@ struct AOTInductorConstantMapEntry { AtenTensorHandle handle; }; +// --------------------------------------------------------------------------- +// C-compatible tensor descriptor for crossing the DSO boundary. +// +// This struct carries the same information as ArrayRefTensor but uses only +// C-compatible types so the host process and DSO can be built with different +// C++ standard libraries (e.g. libc++ vs libstdc++). All pointer fields +// reference memory owned by the caller; no copies are made. +// +// Maximum supported number of dimensions. 8 covers all practical AOTI +// models; tensors with more dims should fall back to the AtenTensorHandle +// interface. +// --------------------------------------------------------------------------- +#define AOTI_ARRAYREF_TENSOR_MAX_DIMS 8 + +struct AOTInductorArrayRefTensor { + // Pointer to the raw data buffer. Not owned. + void* data; + + // Number of elements in the data buffer (product of sizes for contiguous + // tensors). + int64_t numel; + + // Static-size arrays for shape metadata. Only the first `ndim` entries + // are meaningful. + int64_t sizes[AOTI_ARRAYREF_TENSOR_MAX_DIMS]; + int64_t strides[AOTI_ARRAYREF_TENSOR_MAX_DIMS]; + + // Number of dimensions (0 <= ndim <= AOTI_ARRAYREF_TENSOR_MAX_DIMS). + int32_t ndim; + + // Torch dtype encoded as int32_t (same encoding as aoti_torch_dtype_*()). + int32_t dtype; + + // Device information. + int32_t device_type; + int32_t device_idx; + + // Reserved for future extension. Zero-initialize and do not read — a + // newer reader must tolerate zeros, and an older reader must ignore them. + int64_t reserved[4]; +}; + +static_assert( + sizeof(AOTInductorArrayRefTensor) == 192, + "changing the size of AOTInductorArrayRefTensor breaks ABI compatibility!"); + // TODO: Deprecate this API. This was kept for BC compatibility. // Please use AOTInductorModelContainerCreateWithDevice instead. AOTI_API AOTIRuntimeError AOTInductorModelContainerCreate( @@ -265,4 +311,25 @@ AOTI_API AOTIRuntimeError AOTInductorModelContainerGetCallSpec( const char** in_spec, const char** out_spec); +// --------------------------------------------------------------------------- +// C-ABI-safe variant of AOTInductorModelRunMinimalArrayrefInterface. +// +// Instead of passing std::tuple...>& (which encodes C++ +// standard library types into the ABI), this function accepts flat C arrays +// of AOTInductorArrayRefTensor descriptors. The descriptors reference the +// same underlying data buffers -- no copies are made. +// +// The host process marshals its ArrayRefTensor objects into +// AOTInductorArrayRefTensor descriptors, calls into the DSO through this +// pure-C interface, and then unmarshals the output descriptors back. +// Because only C types cross the DSO boundary, the host and DSO can be +// built with different C++ standard libraries (e.g. libc++ vs libstdc++). +// --------------------------------------------------------------------------- +AOTI_API AOTIRuntimeError AOTInductorModelRunMinimalArrayrefInterfaceV2( + AOTInductorModelHandle model_handle, + int32_t num_inputs, + const AOTInductorArrayRefTensor* inputs, + int32_t num_outputs, + AOTInductorArrayRefTensor* outputs); + } // extern "C" diff --git a/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h b/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h index 56bb478396004..b85b69449edc2 100644 --- a/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h +++ b/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h @@ -134,8 +134,13 @@ static std::unique_ptr _createKernel( uint32_t numWarps, uint32_t sharedMemory, void** params, - sycl::queue* queuePtr, - uint32_t threadsPerWarp) { + sycl::queue* queuePtr) { + uint32_t threadsPerWarp = kernelPtr->get_info< + sycl::info::kernel_device_specific::compile_sub_group_size>( + queuePtr->get_device()); + if (threadsPerWarp == 0) { + threadsPerWarp = 32; // default to 32 if not set + } std::string kernelName = kernelPtr->get_info(); uint32_t numParams = kernelPtr->get_info(); diff --git a/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h b/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h index 77072a4e68920..21ee7438165dc 100644 --- a/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h +++ b/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h @@ -121,12 +121,18 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_pow_Tensor_Scalar(AtenTensorHand AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_pow_Tensor_Tensor(AtenTensorHandle self, AtenTensorHandle exponent, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_repeat_interleave_Tensor(AtenTensorHandle repeats, int64_t* output_size, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cpu_replication_pad1d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0); diff --git a/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h b/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h index f047184bb8ca3..413a393f6fa7d 100644 --- a/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h +++ b/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h @@ -135,12 +135,18 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_pow_Tensor_Scalar(AtenTensorHan AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_pow_Tensor_Tensor(AtenTensorHandle self, AtenTensorHandle exponent, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_repeat_interleave_Tensor(AtenTensorHandle repeats, int64_t* output_size, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_cuda_replication_pad1d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0); diff --git a/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h b/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h index 43b2570d1c69b..1a6c153c4a781 100644 --- a/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h +++ b/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h @@ -25,7 +25,7 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__fused_moving_avg_obs_fq_helper_ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__fused_rms_norm(AtenTensorHandle input, const int64_t* normalized_shape, int64_t normalized_shape_len_, AtenTensorHandle* weight, double* eps, AtenTensorHandle* ret0, AtenTensorHandle* ret1); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__grouped_mm(AtenTensorHandle self, AtenTensorHandle mat2, AtenTensorHandle* offs, AtenTensorHandle* bias, int32_t* out_dtype, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__histogramdd_from_bin_cts(AtenTensorHandle self, const int64_t* bins, int64_t bins_len_, const double** range, int64_t range_len_, AtenTensorHandle* weight, int32_t density, AtenTensorHandle* ret0); -AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_attention_math_for_mps(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_mask, double dropout_p, int32_t is_causal, AtenTensorHandle* dropout_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_attention_math_for_mps(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_mask, double dropout_p, int32_t is_causal, AtenTensorHandle* dropout_mask, double* scale, int32_t enable_gqa, AtenTensorHandle* ret0, AtenTensorHandle* ret1); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_fused_attention_overrideable(AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle* attn_bias, double dropout_p, int32_t is_causal, int32_t return_debug_mask, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3, int64_t* ret4, int64_t* ret5, AtenTensorHandle* ret6, AtenTensorHandle* ret7, AtenTensorHandle* ret8); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__scaled_dot_product_fused_attention_overrideable_backward(AtenTensorHandle grad_out, AtenTensorHandle query, AtenTensorHandle key, AtenTensorHandle value, AtenTensorHandle attn_bias, const int32_t* grad_input_mask, int64_t grad_input_mask_len_, AtenTensorHandle out, AtenTensorHandle logsumexp, AtenTensorHandle cum_seq_q, AtenTensorHandle cum_seq_k, int64_t max_q, int64_t max_k, double dropout_p, int32_t is_causal, AtenTensorHandle philox_seed, AtenTensorHandle philox_offset, double* scale, AtenTensorHandle* ret0, AtenTensorHandle* ret1, AtenTensorHandle* ret2, AtenTensorHandle* ret3); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps__to_sparse(AtenTensorHandle self, int32_t* layout, const int64_t** blocksize, int64_t blocksize_len_, int64_t* dense_dim, AtenTensorHandle* ret0); @@ -88,6 +88,7 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_nanmedian(AtenTensorHandle self, AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_narrow(AtenTensorHandle self, int64_t dim, int64_t start, int64_t length, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_native_dropout(AtenTensorHandle input, double p, int32_t* train, AtenTensorHandle* ret0, AtenTensorHandle* ret1); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_nonzero(AtenTensorHandle self, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_nonzero_static(AtenTensorHandle self, int64_t size, int64_t fill_value, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_normal_functional(AtenTensorHandle self, double mean, double std, AtenGeneratorHandle* generator, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_pad(AtenTensorHandle self, const int64_t* pad, int64_t pad_len_, const char* mode, double* value, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_permute(AtenTensorHandle self, const int64_t* dims, int64_t dims_len_, AtenTensorHandle* ret0); @@ -97,12 +98,18 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_pow_Tensor_Scalar(AtenTensorHand AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_pow_Tensor_Tensor(AtenTensorHandle self, AtenTensorHandle exponent, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_repeat_interleave_Tensor(AtenTensorHandle repeats, int64_t* output_size, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_mps_replication_pad1d_backward(AtenTensorHandle grad_output, AtenTensorHandle self, const int64_t* padding, int64_t padding_len_, AtenTensorHandle* ret0); diff --git a/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h b/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h index 307ada8eb8203..ea0cbfdc5659d 100644 --- a/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h +++ b/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h @@ -54,12 +54,18 @@ AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_permute(AtenTensorHandle self, c AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_polar(AtenTensorHandle abs, AtenTensorHandle angle, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_rand_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint(int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_generator(int64_t high, const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_low(int64_t low, int64_t high, const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_low_out(AtenTensorHandle out, int64_t low, int64_t high, const int64_t* size, int64_t size_len_); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_like(AtenTensorHandle self, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randint_like_low_dtype(AtenTensorHandle self, int64_t low, int64_t high, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn(const int64_t* size, int64_t size_len_, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn_generator(const int64_t* size, int64_t size_len_, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn_like(AtenTensorHandle self, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); +AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randn_like_generator(AtenTensorHandle self, AtenGeneratorHandle* generator, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, int32_t* memory_format, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_randperm(int64_t n, int32_t* dtype, int32_t* layout, int32_t* device, int32_t device_index_, int32_t* pin_memory, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_reshape(AtenTensorHandle self, const int64_t* shape, int64_t shape_len_, AtenTensorHandle* ret0); AOTI_TORCH_EXPORT AOTITorchError aoti_torch_xpu_resize_as_(AtenTensorHandle self, AtenTensorHandle the_template, int32_t* memory_format); diff --git a/torch/csrc/inductor/aoti_torch/generated_enum_converters.h b/torch/csrc/inductor/aoti_torch/generated_enum_converters.h index 820023c1031e6..14b97a7f4bc26 100644 --- a/torch/csrc/inductor/aoti_torch/generated_enum_converters.h +++ b/torch/csrc/inductor/aoti_torch/generated_enum_converters.h @@ -50,6 +50,8 @@ inline c10::ScalarType convertSerializedScalarType(int serialized_value) { static_cast(c10::ScalarType::Float8_e4m3fnuz), // 31 static_cast(c10::ScalarType::Float8_e5m2fnuz), // 32 static_cast(c10::ScalarType::Float8_e8m0fnu), // 33 + static_cast(c10::ScalarType::UInt32), // 34 + static_cast(c10::ScalarType::UInt64), // 35 }; constexpr int kMapSize = sizeof(kScalarTypeMap) / sizeof(kScalarTypeMap[0]); diff --git a/torch/csrc/inductor/array_ref_impl.h b/torch/csrc/inductor/array_ref_impl.h index 8cfbc12fb2c3d..e36a3653009ad 100644 --- a/torch/csrc/inductor/array_ref_impl.h +++ b/torch/csrc/inductor/array_ref_impl.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include diff --git a/torch/csrc/inductor/cpp_wrapper/cuda.h b/torch/csrc/inductor/cpp_wrapper/cuda.h index 782a2b677276a..eb9888356c6ff 100644 --- a/torch/csrc/inductor/cpp_wrapper/cuda.h +++ b/torch/csrc/inductor/cpp_wrapper/cuda.h @@ -1,4 +1,10 @@ #pragma once - #include #include +#include + +#ifdef TORCH_INDUCTOR_PRECOMPILE_HEADERS +#include +#include +#include +#endif diff --git a/torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h b/torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h new file mode 100644 index 0000000000000..27f6e81198385 --- /dev/null +++ b/torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h @@ -0,0 +1,201 @@ +#pragma once + +#include + +#include +#include +#if defined(USE_XPU) +#include +#else +#include +#endif + +struct LazyKernelCompileResult { + std::string cubin_path; + std::string mangled_name; + int num_warps; + int shared_mem; + int xblock; + int yblock; + int zblock; + int r0block; + int rsplit; + int rsplit_size; + int config_index; + int global_scratch; + int profile_scratch; +}; + +static PyObject* (*_THPVariable_Wrap)(const at::TensorBase&) = nullptr; +static int32_t (*_THPUtils_unpackInt)(PyObject*) = nullptr; + +// Cached module and function references +static PyObject* triton_lazy_compile_module = nullptr; +static PyObject* start_kernel_compile = nullptr; +static PyObject* run_triton_kernel_with_autotune = nullptr; + +// Per-module dict for pending kernel compile results (avoids global state +// collisions when multiple compiled modules produce kernels with the same +// name). +static PyObject* _module_pending_kernels = nullptr; + +static inline void loadLazyCompileFuncs() { + if (triton_lazy_compile_module == nullptr) { + triton_lazy_compile_module = + PyImport_ImportModule("torch._inductor.runtime.triton_lazy_compile"); + AOTI_TORCH_CHECK( + triton_lazy_compile_module, "Failed to import triton_lazy_compile"); + + start_kernel_compile = PyObject_GetAttrString( + triton_lazy_compile_module, "start_kernel_compile"); + AOTI_TORCH_CHECK( + start_kernel_compile, "Failed to get start_kernel_compile"); + + run_triton_kernel_with_autotune = PyObject_GetAttrString( + triton_lazy_compile_module, "run_triton_kernel_with_autotune"); + AOTI_TORCH_CHECK( + run_triton_kernel_with_autotune, + "Failed to get run_triton_kernel_with_autotune"); + + RAIIPyObject guards_mod = PyImport_ImportModule("torch._C._dynamo.guards"); + AOTI_TORCH_CHECK(guards_mod, "Failed to import torch._C._dynamo.guards"); + + RAIIPyObject wrap_addr = + PyObject_GetAttrString(guards_mod, "_torchinductor_thp_variable_wrap"); + AOTI_TORCH_CHECK( + wrap_addr, "Failed to get _torchinductor_thp_variable_wrap"); + _THPVariable_Wrap = reinterpret_cast( + PyLong_AsVoidPtr(wrap_addr)); + AOTI_TORCH_CHECK(_THPVariable_Wrap, "THPVariable_Wrap not resolved"); + + RAIIPyObject unpack_addr = PyObject_GetAttrString( + guards_mod, "_torchinductor_thputils_unpack_int"); + AOTI_TORCH_CHECK( + unpack_addr, "Failed to get _torchinductor_thputils_unpack_int"); + _THPUtils_unpackInt = reinterpret_cast( + PyLong_AsVoidPtr(unpack_addr)); + AOTI_TORCH_CHECK(_THPUtils_unpackInt, "THPUtils_unpackInt not resolved"); + } +} + +static inline std::string getStringAttr(PyObject* obj, const char* attr) { + RAIIPyObject val = PyObject_GetAttrString(obj, attr); + AOTI_TORCH_CHECK(val, "Failed to get attribute"); + return PyUnicode_AsUTF8(val); +} + +static inline int getIntAttr(PyObject* obj, const char* attr) { + RAIIPyObject val = PyObject_GetAttrString(obj, attr); + AOTI_TORCH_CHECK(val, "Failed to get attribute"); + return _THPUtils_unpackInt(val); +} + +static inline int getOptionalIntAttr( + PyObject* obj, + const char* attr, + int sentinel = -1) { + RAIIPyObject val = PyObject_GetAttrString(obj, attr); + AOTI_TORCH_CHECK(val, "Failed to get attribute"); + return (val.get() != Py_None) ? _THPUtils_unpackInt(val) : sentinel; +} + +static inline LazyKernelCompileResult extractCompileResult(PyObject* result) { + LazyKernelCompileResult compile_result; + compile_result.cubin_path = getStringAttr(result, "cubin_path"); + compile_result.mangled_name = getStringAttr(result, "mangled_name"); + compile_result.num_warps = getIntAttr(result, "num_warps"); + compile_result.shared_mem = getIntAttr(result, "shared_mem"); + compile_result.xblock = getIntAttr(result, "xblock"); + compile_result.yblock = getIntAttr(result, "yblock"); + compile_result.zblock = getIntAttr(result, "zblock"); + compile_result.r0block = getIntAttr(result, "r0block"); + compile_result.rsplit = getIntAttr(result, "rsplit"); + compile_result.rsplit_size = getIntAttr(result, "rsplit_size"); + compile_result.config_index = getOptionalIntAttr(result, "config_index"); + compile_result.global_scratch = getOptionalIntAttr(result, "global_scratch"); + compile_result.profile_scratch = + getOptionalIntAttr(result, "profile_scratch"); + return compile_result; +} + +template +static inline PyObject* convertArgToPython(const T& arg) { + using DecayedT = std::decay_t; + if constexpr (std::is_same_v) { + at::Tensor* tensor_ptr = + torch::aot_inductor::tensor_handle_to_tensor_pointer(arg); + return _THPVariable_Wrap(*tensor_ptr); + } else if constexpr (std::is_same_v< + DecayedT, + torch::aot_inductor::RAIIAtenTensorHandle>) { + at::Tensor* tensor_ptr = + torch::aot_inductor::tensor_handle_to_tensor_pointer(arg.get()); + return _THPVariable_Wrap(*tensor_ptr); + } else if constexpr (std::is_same_v) { + PyObject* py_arg = arg ? Py_True : Py_False; + Py_INCREF(py_arg); + return py_arg; + } else if constexpr (std::is_integral_v) { + return PyLong_FromLongLong(static_cast(arg)); + } else if constexpr (std::is_floating_point_v) { + return PyFloat_FromDouble(static_cast(arg)); + } else { + AOTI_TORCH_CHECK(false, "Invalid input type to convertArgToPython"); + } +} + +template +static inline LazyKernelCompileResult runTritonKernelWithAutotune( + PyObject* pending_kernels, + const std::string& kernel_name, + void* stream, + const Args&... kernel_args) { + py::gil_scoped_acquire_simple acquire; + + constexpr size_t num_args = sizeof...(Args); + RAIIPyObject py_args_list = PyList_New(num_args); + AOTI_TORCH_CHECK(py_args_list, "Failed to create args list"); + + size_t idx = 0; + auto add_arg = [&py_args_list, &idx](PyObject* py_arg) { + AOTI_TORCH_CHECK(py_arg, "Failed to convert argument"); + PyList_SetItem(py_args_list, idx++, py_arg); + }; + // Use array pack-expansion instead of a fold expression to avoid + // hitting the compiler's expression-nesting limit when there are + // hundreds of kernel arguments (e.g. combo kernels). + int dummy[] = {0, (add_arg(convertArgToPython(kernel_args)), 0)...}; + (void)dummy; + + RAIIPyObject call_args = PyTuple_Pack( + 4, + pending_kernels, + PyUnicode_FromString(kernel_name.c_str()), + PyLong_FromVoidPtr(stream), + py_args_list.get()); + AOTI_TORCH_CHECK(call_args, "Failed to create call args"); + + RAIIPyObject result = + PyObject_CallObject(run_triton_kernel_with_autotune, call_args); + AOTI_TORCH_CHECK(result, "Failed to run kernel with autotuning"); + + return extractCompileResult(result); +} + +static inline void startKernelCompile( + PyObject* pending_kernels, + const std::string& kernel_name, + const std::string& kernel_source) { + py::gil_scoped_acquire_simple acquire; + + RAIIPyObject py_name = PyUnicode_FromString(kernel_name.c_str()); + RAIIPyObject py_source = PyUnicode_FromString(kernel_source.c_str()); + AOTI_TORCH_CHECK(py_name && py_source, "Failed to create Python args"); + + RAIIPyObject call_args = + PyTuple_Pack(3, pending_kernels, py_name.get(), py_source.get()); + AOTI_TORCH_CHECK(call_args, "Failed to create call args"); + + RAIIPyObject result = PyObject_CallObject(start_kernel_compile, call_args); + AOTI_TORCH_CHECK(result, "Failed to start kernel compilation"); +} diff --git a/torch/csrc/inductor/inductor_ops.cpp b/torch/csrc/inductor/inductor_ops.cpp index 9723e27e6ba8a..96a41a78c00bc 100644 --- a/torch/csrc/inductor/inductor_ops.cpp +++ b/torch/csrc/inductor/inductor_ops.cpp @@ -108,4 +108,11 @@ TORCH_LIBRARY_FRAGMENT(inductor, m) { {at::Tag::pt2_compliant_tag}); } +TORCH_LIBRARY_FRAGMENT(inductor_prims, m) { + m.def( + "inductor_reserve_rng_state(Generator? generator, SymInt increment) " + "-> (Tensor, Tensor, Tensor)", + {at::Tag::pt2_compliant_tag}); +} + } // namespace torch::inductor diff --git a/torch/csrc/inductor/inductor_ops_gpu.cpp b/torch/csrc/inductor/inductor_ops_gpu.cpp new file mode 100644 index 0000000000000..2133dbcb4f796 --- /dev/null +++ b/torch/csrc/inductor/inductor_ops_gpu.cpp @@ -0,0 +1,89 @@ +#include +#include +#include +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +#include +#include +#include +#endif + +#if defined(USE_CUDA) || defined(USE_ROCM) +#include +#endif + +namespace torch::inductor { +using namespace at; + +#if defined(USE_CUDA) || defined(USE_ROCM) + +// Reserves RNG state for Inductor with CUDA Graph support. +// +// This function allows Inductor to reserve a specific amount of RNG offset +// (increment) for a kernel. It is designed to be safe for CUDA Graph capture +// by explicitly handling the internal generator state via public APIs. +// +// Behavior: +// - Graph Mode: Advances the generator state and returns pointers (wrapped as +// tensors) to the extragraph state. These tensors effectively point to the +// GPU memory that will be updated by `replay_prologue`. +// - Eager Mode: Advances the generator state and returns concrete values +// wrapped in 1D tensors to maintain shape consistency. +// +// -param gen The CUDA generator to use. +// -param increment The number of RNG values to reserve. +// -return A tuple of (Seed Tensor, Offset Tensor, Intragraph Offset CPU +// Tensor). +static std::tuple inductor_reserve_rng_state_impl( + const Generator& generator, + c10::SymInt increment) { + auto* gen_impl = at::check_generator(generator); + + const auto dev_opts = + at::TensorOptions().dtype(at::kLong).device(generator.device()); + const auto cpu_opts = at::TensorOptions().dtype(at::kLong).device(at::kCPU); + + int64_t inc = increment.expect_int(); + const at::PhiloxCudaState st = + gen_impl->philox_cuda_state(static_cast(inc)); + + if (st.captured_) { + auto seed_t = at::from_blob( + static_cast(st.seed_.ptr), {1}, [](void*) {}, dev_opts); + auto off_t = at::from_blob( + static_cast(st.offset_.ptr), {1}, [](void*) {}, dev_opts); + auto intra_t = + at::scalar_tensor(static_cast(st.offset_intragraph_), cpu_opts) + .unsqueeze(0); + return {seed_t, off_t, intra_t}; + } + + auto seed_t = at::scalar_tensor(static_cast(st.seed_.val), dev_opts) + .unsqueeze(0); + auto off_t = at::scalar_tensor(static_cast(st.offset_.val), dev_opts) + .unsqueeze(0); + auto intra_t = at::zeros({1}, cpu_opts); + return {seed_t, off_t, intra_t}; +} + +TORCH_LIBRARY_IMPL(inductor_prims, BackendSelect, m) { + m.impl( + "inductor_reserve_rng_state", TORCH_FN(inductor_reserve_rng_state_impl)); +} + +TORCH_LIBRARY_IMPL(inductor_prims, CUDA, m) { + m.impl( + "inductor_reserve_rng_state", TORCH_FN(inductor_reserve_rng_state_impl)); +} + +TORCH_LIBRARY_IMPL(inductor_prims, HIP, m) { + m.impl( + "inductor_reserve_rng_state", TORCH_FN(inductor_reserve_rng_state_impl)); +} + +#endif + +} // namespace torch::inductor diff --git a/torch/csrc/jit/frontend/schema_type_parser.cpp b/torch/csrc/jit/frontend/schema_type_parser.cpp index 1884ac61f8c96..960fc9c260666 100644 --- a/torch/csrc/jit/frontend/schema_type_parser.cpp +++ b/torch/csrc/jit/frontend/schema_type_parser.cpp @@ -108,6 +108,9 @@ TypePtr SchemaTypeParser::parseBaseType() { {"Any", c10::TypeFactory::get()}, {"AnyClassType", c10::TypeFactory::get()}, {"AnyEnumType", c10::TypeFactory::get()}, + // PyObjectType::get() used directly because PyObjectType is excluded + // from FORALL_DYNAMIC_TYPES (not supported on xplat/mobile) + {"PyObject", c10::PyObjectType::get()}, }; auto tok = L.cur(); if (!L.nextIf(TK_NONE) && !L.nextIf(TK_NONE_TYPE)) { diff --git a/torch/csrc/jit/ir/ir.cpp b/torch/csrc/jit/ir/ir.cpp index 9b85b0909826f..4d1a4594c79a6 100644 --- a/torch/csrc/jit/ir/ir.cpp +++ b/torch/csrc/jit/ir/ir.cpp @@ -929,12 +929,7 @@ void Value::replaceAllUsesAfterNodeWith(const Node* node, Value* newValue) { } }); - uses_.erase( - std::remove_if( - uses_.begin(), - uses_.end(), - [&node](const Use& u) { return u.user->isAfter(node); }), - uses_.end()); + std::erase_if(uses_, [&node](const Use& u) { return u.user->isAfter(node); }); } void Value::replaceAllUsesDominatedByNodeWith( @@ -947,12 +942,8 @@ void Value::replaceAllUsesDominatedByNodeWith( } }); - uses_.erase( - std::remove_if( - uses_.begin(), - uses_.end(), - [&node](const Use& u) { return u.user->isDominatedBy(node); }), - uses_.end()); + std::erase_if( + uses_, [&node](const Use& u) { return u.user->isDominatedBy(node); }); } static size_t findArgument( diff --git a/torch/csrc/jit/passes/onnx/constant_map.cpp b/torch/csrc/jit/passes/onnx/constant_map.cpp index 60699a1e75ef4..c332ea4204eb9 100644 --- a/torch/csrc/jit/passes/onnx/constant_map.cpp +++ b/torch/csrc/jit/passes/onnx/constant_map.cpp @@ -120,28 +120,25 @@ std::optional> ConstantValueMap:: GetShapeInto1DInt64VectorWithOneUnknown(const std::string& value_name) { if (ConstantValueMap::HasShape(value_name)) { auto shape_size = ConstantValueMap::GetShape(value_name).value(); - std::vector shape_value; if (shape_size.isComplete()) { - shape_value = - ConstantValueMap::GetCompleteShapeInto1DInt64Vector(shape_size); - return shape_value; - } else { - size_t count_unknown = 0; - auto shape_size_sizes = shape_size.sizes(); - if (shape_size_sizes.has_value()) { - auto shape_symbol_list = shape_size_sizes.value(); - for (const auto& v : shape_symbol_list) { - if (v.is_static()) { - shape_value.emplace_back(v.static_size()); - } else { - shape_value.emplace_back(-1); - count_unknown += 1; - } - } - if (count_unknown == 1) { - return shape_value; + return ConstantValueMap::GetCompleteShapeInto1DInt64Vector(shape_size); + } + size_t count_unknown = 0; + auto shape_size_sizes = shape_size.sizes(); + if (shape_size_sizes.has_value()) { + std::vector shape_value; + auto shape_symbol_list = shape_size_sizes.value(); + for (const auto& v : shape_symbol_list) { + if (v.is_static()) { + shape_value.emplace_back(v.static_size()); + } else { + shape_value.emplace_back(-1); + count_unknown += 1; } } + if (count_unknown == 1) { + return shape_value; + } } } return std::nullopt; diff --git a/torch/csrc/jit/passes/onnx/function_extraction.cpp b/torch/csrc/jit/passes/onnx/function_extraction.cpp index c82fd3d63c9af..eee0251f4153b 100644 --- a/torch/csrc/jit/passes/onnx/function_extraction.cpp +++ b/torch/csrc/jit/passes/onnx/function_extraction.cpp @@ -710,14 +710,8 @@ void FunctionExtractor::ConvertScopeToFunction( ctx_nlist.insert(last_n_it, func_n); // remove replaced nodes from list - ctx_nlist.erase( - std::remove_if( - ctx_nlist.begin(), - ctx_nlist.end(), - [&old_nodes](Node* n) { - return old_nodes.find(n) != old_nodes.end(); - }), - ctx_nlist.end()); + std::erase_if( + ctx_nlist, [&old_nodes](Node* n) { return old_nodes.contains(n); }); GRAPH_DEBUG("Parent total nodes after remove: ", ctx_nlist.size()); diff --git a/torch/csrc/jit/passes/onnx/helper.cpp b/torch/csrc/jit/passes/onnx/helper.cpp index 3897a8d5cae5e..a3bcb9c49f8e2 100644 --- a/torch/csrc/jit/passes/onnx/helper.cpp +++ b/torch/csrc/jit/passes/onnx/helper.cpp @@ -39,14 +39,8 @@ void eraseUnusedBlockInputs(Block* b) { } void eraseUnusedValuesFromMap(ValueToParamPairMap& valsToParamsMap) { - auto it = valsToParamsMap.begin(); - while (it != valsToParamsMap.end()) { - if (!it->first->hasUses()) { - it = valsToParamsMap.erase(it); - } else { - ++it; - } - } + std::erase_if( + valsToParamsMap, [](const auto& pr) { return !pr.first->hasUses(); }); } void buildParamsMapFromValueToParamsMap( diff --git a/torch/csrc/jit/python/init.cpp b/torch/csrc/jit/python/init.cpp index 931d80ffd6a68..4587d70b1e0ac 100644 --- a/torch/csrc/jit/python/init.cpp +++ b/torch/csrc/jit/python/init.cpp @@ -1796,7 +1796,7 @@ void initJITBindings(PyObject* module) { m.def( "_jit_get_operation", - [](const std::string& op_name) { + [](const std::string& op_name) -> py::tuple { try { auto symbol = Symbol::fromQualString(op_name); const auto sortedOps = getAllSortedOperatorsFor(symbol); @@ -1843,7 +1843,7 @@ void initJITBindings(PyObject* module) { "_maybe_call_torch_function_for_op_packet", [](py::handle op_overload_packet, const py::args& args, - const py::kwargs& kwargs) { + const py::kwargs& kwargs) -> py::tuple { py::list ns_method = op_overload_packet.attr("_qualified_op_name").attr("split")("::"); auto res = _maybe_handle_torch_function( diff --git a/torch/csrc/jit/python/pybind_utils.cpp b/torch/csrc/jit/python/pybind_utils.cpp index cf422102f7dec..bbcc31515ce3e 100644 --- a/torch/csrc/jit/python/pybind_utils.cpp +++ b/torch/csrc/jit/python/pybind_utils.cpp @@ -543,9 +543,15 @@ IValue toIValue(py::handle obj, const TypePtr& type, std::optional N) { } case TypeKind::CapsuleType: { #ifdef USE_DISTRIBUTED - // Handle ProcessGroup custom class as a capsule + // Handle ProcessGroup custom class as a capsule. FakeScriptObject + // (used during Dynamo tracing with CooR) passes py::isinstance via + // OpaqueBaseMeta but cannot be cast directly; unwrap real_obj first. if (py::isinstance(obj)) { - auto cpp_obj = obj.cast>(); + py::handle target = obj; + if (py::hasattr(obj, "real_obj")) { + target = obj.attr("real_obj"); + } + auto cpp_obj = target.cast>(); return IValue::make_capsule(cpp_obj); } #endif @@ -1006,6 +1012,18 @@ std::optional detail::_tryToInferTypeImpl(py::handle input) { if (py::isinstance(input)) { return InferredType(CapsuleType::get()); } + // During Dynamo tracing with compile-on-one-rank (CooR), opaque reference + // types like ProcessGroup are wrapped in FakeScriptObject. Python-level + // isinstance() sees through the wrapper (via OpaqueBaseMeta), but the C++ + // py::isinstance above does too — yet the subsequent pybind11 cast would + // fail because FakeScriptObject is not a C++ bound object. Detect this + // case by checking for the wrapped real_obj attribute. + if (py::hasattr(input, "real_obj")) { + py::object real = input.attr("real_obj"); + if (py::isinstance(real)) { + return InferredType(CapsuleType::get()); + } + } #endif return std::nullopt; diff --git a/torch/csrc/jit/python/pybind_utils.h b/torch/csrc/jit/python/pybind_utils.h index 378bf4636fd2e..37e49c2952620 100644 --- a/torch/csrc/jit/python/pybind_utils.h +++ b/torch/csrc/jit/python/pybind_utils.h @@ -23,7 +23,6 @@ #include #include #include -#include #ifdef USE_DISTRIBUTED #include #include @@ -429,6 +428,14 @@ inline InferredType tryToInferType(py::handle input) { return InferredType(IntType::get()); } + // Check for types registered in _tryToInferTypeImpl (e.g. ProcessGroup) + // before falling through to the expensive inspect.isclass / JIT compilation + // path below. + auto ty = detail::_tryToInferTypeImpl(input); + if (ty.has_value()) { + return ty.value(); + } + auto enum_type = py::module::import("enum").attr("Enum"); py::bool_ isEnumValue = py::isinstance(input, enum_type); if (py::cast(isEnumValue)) { @@ -518,11 +525,6 @@ inline InferredType tryToInferType(py::handle input) { return InferredType("Cannot infer concrete type of torch.nn.Module"); } - auto ty = detail::_tryToInferTypeImpl(input); - if (ty.has_value()) { - return ty.value(); - } - // Try container types return tryToInferContainerType(input, false); } @@ -554,7 +556,7 @@ inline InferredType tryToInferPrimitiveType(py::handle input) { inline InferredType tryToInferContainerType( py::handle input, bool primitiveTypeOnly = false) { - if (six::isTuple(input)) { + if (PyTuple_Check(input.ptr())) { py::tuple tuple = py::cast(input); std::vector element_types; element_types.reserve(tuple.size()); diff --git a/torch/csrc/jit/python/python_arg_flatten.cpp b/torch/csrc/jit/python/python_arg_flatten.cpp index b71f21d043a31..fde4b35f37560 100644 --- a/torch/csrc/jit/python/python_arg_flatten.cpp +++ b/torch/csrc/jit/python/python_arg_flatten.cpp @@ -1,7 +1,7 @@ #include #include #include -#include +#include #include @@ -44,7 +44,7 @@ py::object cast_handle_sequence(std::vector objs) { void flatten_rec(PyObject* obj, ParsedArgs& args) { auto& structure = args.desc.structure; - if (six::isTuple(obj)) { + if (PyTuple_Check(obj)) { structure.push_back(D::TupleOpen); for (auto item : py::reinterpret_borrow(obj)) flatten_rec(item.ptr(), args); diff --git a/torch/csrc/jit/runtime/graph_executor.cpp b/torch/csrc/jit/runtime/graph_executor.cpp index 4bdab8c5dcb22..fa37ccab96d06 100644 --- a/torch/csrc/jit/runtime/graph_executor.cpp +++ b/torch/csrc/jit/runtime/graph_executor.cpp @@ -635,6 +635,9 @@ struct GraphExecutorImpl : public GraphExecutorImplBase { const ExecutionPlan& getPlanFor( Stack& stack, std::optional remaining_bailout_depth) override { + if (FLAGS_torch_jit_input_independent_optimization) { + return getInputIndependentPlan(); + } return getGraphExecutorOptimize() ? getOrCompile(stack) : getOrCompileFallback(); } @@ -778,6 +781,38 @@ struct GraphExecutorImpl : public GraphExecutorImplBase { return ExecutionPlan(opt_graph, function_name_); } + const ExecutionPlan& getInputIndependentPlan() override { + std::lock_guard lock(compile_mutex); + if (!input_independent_plan_) { + auto opt_graph = graph->copy(); + Inline(*opt_graph); + LowerGradOf(*opt_graph); + specializeAutogradZero(opt_graph); + LowerSimpleTuples(opt_graph); + ConstantPooling(opt_graph); + runRequiredPasses(opt_graph); + ConstantPropagation(opt_graph); + // Skip PropagateInputShapes and PropagateRequiresGrad since they need + // actual input data. + runOptimization(opt_graph); + + // Input-independent passes from runNondiffOptimization. Skipped: + // FuseTensorExprs/FuseGraph (need specialized tensor types). + for (const auto& passPair : getCustomPrePasses()) { + passPair.first(opt_graph); + } + DecomposeOps(opt_graph); + BatchMM(opt_graph); + for (const auto& passPair : getCustomPostPasses()) { + passPair.first(opt_graph); + } + + EliminateDeadCode(opt_graph); + input_independent_plan_ = ExecutionPlan(opt_graph, function_name_); + } + return *input_independent_plan_; + } + ~GraphExecutorImpl() override = default; // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes) @@ -787,6 +822,10 @@ struct GraphExecutorImpl : public GraphExecutorImplBase { // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes) ExecutionPlan fallback; + // Cached plan from getOptimizedPlan() -- uses only input-independent passes. + // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes) + std::optional input_independent_plan_; + // Mapping from argument configurations to optimized versions of the graph // that are specialized to the spec. // NOLINTNEXTLINE(cppcoreguidelines-non-private-member-variables-in-classes) @@ -842,6 +881,10 @@ const ExecutionPlan& GraphExecutor::getPlanFor( return pImpl->getPlanFor(inputs, remaining_bailout_depth); } +const ExecutionPlan& GraphExecutor::getInputIndependentPlan() { + return pImpl->getInputIndependentPlan(); +} + GraphExecutorState GraphExecutor::getDebugState() { return pImpl->getDebugState(); } diff --git a/torch/csrc/jit/runtime/graph_executor.h b/torch/csrc/jit/runtime/graph_executor.h index d1039216de3ea..94ed97f1784e8 100644 --- a/torch/csrc/jit/runtime/graph_executor.h +++ b/torch/csrc/jit/runtime/graph_executor.h @@ -87,6 +87,11 @@ struct TORCH_API GraphExecutor { const ExecutionPlan& getPlanFor( Stack& inputs, std::optional remaining_bailout_depth = std::nullopt); + // Returns an optimized execution plan without requiring input arguments. + // Runs input-independent optimization passes (e.g. inlining, constant + // propagation, peephole, CSE) but skips profiling-based specializations + // that require runtime type/shape information. + const ExecutionPlan& getInputIndependentPlan(); GraphExecutorState getDebugState(); void debugFlushCompilationCache(); diff --git a/torch/csrc/jit/runtime/graph_executor_impl.h b/torch/csrc/jit/runtime/graph_executor_impl.h index 70069ac1907b0..78c7a7b79d982 100644 --- a/torch/csrc/jit/runtime/graph_executor_impl.h +++ b/torch/csrc/jit/runtime/graph_executor_impl.h @@ -79,6 +79,11 @@ struct GraphExecutorImplBase { virtual const ExecutionPlan& getPlanFor( Stack& stack, std::optional remaining_bailout_depth = std::nullopt) = 0; + // Returns an optimized execution plan without requiring input arguments. + // Runs input-independent optimization passes (e.g. inlining, constant + // propagation, peephole, CSE) but skips profiling-based specializations + // that require runtime type/shape information. + virtual const ExecutionPlan& getInputIndependentPlan() = 0; virtual GraphExecutorState getDebugState() = 0; virtual ~GraphExecutorImplBase() = default; diff --git a/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp b/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp index 680244b363c36..c638c7f80a39b 100644 --- a/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp +++ b/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp @@ -56,6 +56,13 @@ C10_DEFINE_bool( false, "fuse on 12 dynamic compilations") +C10_DEFINE_bool( + torch_jit_input_independent_optimization, + false, + "If set, getPlanFor will use input-independent optimization passes only, " + "skipping profiling-based specializations that require runtime type/shape " + "information. Useful for predictor nets and AOT compilation scenarios.") + C10_DEFINE_bool( torch_jit_release_profiling_graph_after_optimization, false, @@ -714,10 +721,59 @@ const ExecutionPlan& ProfilingGraphExecutorImpl::getPlanFor( } return *optimized_plan_; } + if (FLAGS_torch_jit_input_independent_optimization) { + return getInputIndependentPlanImpl(); + } // if depth is not set, use return getOptimizedPlanFor(stack, remaining_bailout_depth); } +const ExecutionPlan& ProfilingGraphExecutorImpl::getInputIndependentPlan() { + std::lock_guard lock(compile_mutex); + return getInputIndependentPlanImpl(); +} + +const ExecutionPlan& +ProfilingGraphExecutorImpl::getInputIndependentPlanImpl() { + if (optimized_plan_) { + return *optimized_plan_; + } + + auto copy = graph->copy(); + if (!getGraphExecutorOptimize() || !getProfilingMode()) { + LowerGradOf(*copy); + RemoveExpands(copy); + } else { + // Run all input-independent optimizations. This includes + // runProfilingInsensitiveOptimizations (inlining, constant propagation, + // CSE, peephole, etc.) followed by runPreAutodiffPassPipeline which adds + // loop unrolling, list mutation removal, and additional rounds of + // peephole + constant propagation. Profiling-dependent passes (type + // specialization, fusion, autodiff guards) are skipped since they + // require runtime type/shape information from actual inputs. + runProfilingInsensitiveOptimizations(copy); + runPreAutodiffPassPipeline(copy); + + // Run additional input-independent passes from runNoGradOptimizations + // and runFinalOptimizations. These operate on graph structure (node + // patterns, alias analysis) and do not require profiled type/shape info. + // Skipped: RemoveProfileNodesAndSpecializeTypes, FuseTensorExprs, + // FuseGraph (these need profiled tensor types for specialization/fusion). + for (const auto& passPair : getCustomPrePasses()) { + passPair.first(copy); + } + BatchMM(copy); + for (const auto& passPair : getCustomPostPasses()) { + passPair.first(copy); + } + AddIfThenElseOp(copy); + EliminateDeadCode(copy); + } + optimized_plan_ = ExecutionPlan(copy, function_name_); + time_optimized_plan_created_ = getNowInSecs(); + return *optimized_plan_; +} + GraphExecutorState ProfilingGraphExecutorImpl::getDebugState() { GraphExecutorState state; TORCH_INTERNAL_ASSERT(optimized_plan_); diff --git a/torch/csrc/jit/runtime/profiling_graph_executor_impl.h b/torch/csrc/jit/runtime/profiling_graph_executor_impl.h index da437bb456e92..422a7646cef8a 100644 --- a/torch/csrc/jit/runtime/profiling_graph_executor_impl.h +++ b/torch/csrc/jit/runtime/profiling_graph_executor_impl.h @@ -7,6 +7,7 @@ TORCH_DECLARE_bool(torch_jit_static_then_dynamic); TORCH_DECLARE_bool(torch_jit_always_dynamic); +C10_DECLARE_bool(torch_jit_input_independent_optimization); C10_DECLARE_bool(torch_jit_release_profiling_graph_after_optimization); C10_DECLARE_int32(torch_jit_release_profiling_graph_delay_in_seconds); C10_DECLARE_int64(torch_jit_num_profiled_runs); @@ -24,6 +25,7 @@ struct TORCH_API ProfilingGraphExecutorImpl : public GraphExecutorImplBase { const ExecutionPlan& getPlanFor( Stack& stack, std::optional remaining_bailout_depth) override; + const ExecutionPlan& getInputIndependentPlan() override; GraphExecutorState getDebugState() override; ~ProfilingGraphExecutorImpl() override = default; @@ -37,6 +39,8 @@ struct TORCH_API ProfilingGraphExecutorImpl : public GraphExecutorImplBase { const ExecutionPlan& getOptimizedPlanFor( Stack& stack, std::optional remaining_bailout_depth); + // Input-independent optimization, assumes compile_mutex is held. + const ExecutionPlan& getInputIndependentPlanImpl(); void runProfilingInsensitiveOptimizations(std::shared_ptr& graph); void runProfilingOptimizations( std::shared_ptr& graph, diff --git a/torch/csrc/jit/runtime/register_prim_ops.cpp b/torch/csrc/jit/runtime/register_prim_ops.cpp index 1f168d24e8adf..e512853a57567 100644 --- a/torch/csrc/jit/runtime/register_prim_ops.cpp +++ b/torch/csrc/jit/runtime/register_prim_ops.cpp @@ -583,56 +583,6 @@ static const std::vector opGenArgs{ push(stack, arg.element_size()); }, aliasAnalysisFromSchema()), - OperatorGeneratorArgs( - TORCH_SELECTIVE_SCHEMA("aten::numel(Tensor self) -> int"), - [](Stack& stack) { - at::Tensor arg = pop(stack).toTensor(); - push(stack, arg.numel()); - }, - aliasAnalysisFromSchema()), - OperatorGeneratorArgs( - TORCH_SELECTIVE_SCHEMA("aten::dim(Tensor self) -> int"), - dim, - aliasAnalysisFromSchema()), - OperatorGeneratorArgs( - TORCH_SELECTIVE_SCHEMA("aten::get_device(Tensor self) -> int"), - [](Stack& stack) { - RECORD_FUNCTION("get_device", c10::ArrayRef{}); - auto result = - at::get_device((std::move(peek(stack, 0, 1))).toTensor()); - drop(stack, 1); - pack(stack, result); - }, - aliasAnalysisFromSchema()), - OperatorGeneratorArgs( - TORCH_SELECTIVE_SCHEMA("aten::storage_offset(Tensor self) -> int"), - [](Stack& stack) { - RECORD_FUNCTION("storage_offset", c10::ArrayRef{}); - auto result = - ((std::move(peek(stack, 0, 1))).toTensor()).storage_offset(); - drop(stack, 1); - pack(stack, result); - }, - aliasAnalysisFromSchema()), - OperatorGeneratorArgs( - TORCH_SELECTIVE_SCHEMA("aten::is_contiguous(Tensor self) -> bool"), - [](Stack& stack) { - RECORD_FUNCTION("is_contiguous", c10::ArrayRef{}); - auto result = - ((std::move(peek(stack, 0, 1))).toTensor()).is_contiguous(); - drop(stack, 1); - pack(stack, result); - }, - aliasAnalysisFromSchema()), - OperatorGeneratorArgs( - TORCH_SELECTIVE_SCHEMA( - "aten::is_contiguous.memory_format(Tensor self, MemoryFormat memory_format) -> bool"), - [](Stack& stack) { - auto memory_format = pop(stack).toMemoryFormat(); - auto t = pop(stack).toTensor(); - push(stack, t.is_contiguous(memory_format)); - }, - aliasAnalysisFromSchema()), OperatorGeneratorArgs( // NB: intentionally suffixed with extra _format to prevent tests for // "_like" suffix from triggering on this diff --git a/torch/csrc/jit/runtime/simple_graph_executor_impl.cpp b/torch/csrc/jit/runtime/simple_graph_executor_impl.cpp index fd908b48ee043..4671ab0fdcaf7 100644 --- a/torch/csrc/jit/runtime/simple_graph_executor_impl.cpp +++ b/torch/csrc/jit/runtime/simple_graph_executor_impl.cpp @@ -28,6 +28,22 @@ const ExecutionPlan& SimpleGraphExecutorImpl::getPlanFor( return *execution_plan_; } +const ExecutionPlan& SimpleGraphExecutorImpl::getInputIndependentPlan() { + std::lock_guard lock(compile_mutex); + return getInputIndependentPlanImpl(); +} + +const ExecutionPlan& SimpleGraphExecutorImpl::getInputIndependentPlanImpl() { + if (execution_plan_) { + return *execution_plan_; + } + auto copy = graph->copy(); + runNooptPassPipeline(copy); + execution_plan_ = ExecutionPlan(copy, function_name_); + + return *execution_plan_; +} + GraphExecutorState SimpleGraphExecutorImpl::getDebugState() { GraphExecutorState state; TORCH_INTERNAL_ASSERT(execution_plan_); diff --git a/torch/csrc/jit/runtime/simple_graph_executor_impl.h b/torch/csrc/jit/runtime/simple_graph_executor_impl.h index e1ebed46ede80..d972432b6883e 100644 --- a/torch/csrc/jit/runtime/simple_graph_executor_impl.h +++ b/torch/csrc/jit/runtime/simple_graph_executor_impl.h @@ -13,10 +13,13 @@ struct TORCH_API SimpleGraphExecutorImpl : public GraphExecutorImplBase { const ExecutionPlan& getPlanFor( Stack& stack, std::optional remaining_bailout_depth) override; + const ExecutionPlan& getInputIndependentPlan() override; GraphExecutorState getDebugState() override; ~SimpleGraphExecutorImpl() override = default; private: + const ExecutionPlan& getInputIndependentPlanImpl(); + std::optional execution_plan_; }; diff --git a/torch/csrc/jit/tensorexpr/ir_simplifier.cpp b/torch/csrc/jit/tensorexpr/ir_simplifier.cpp index 53eca369cc5ac..74745ed39a40c 100644 --- a/torch/csrc/jit/tensorexpr/ir_simplifier.cpp +++ b/torch/csrc/jit/tensorexpr/ir_simplifier.cpp @@ -37,7 +37,7 @@ static inline ExprPtr newBinaryOpOfType( case IRNodeType::kRshift: return alloc(lhs, rhs); default: - LOG(FATAL) << "unsupported expr_type: " << static_cast(expr_type); + TORCH_INTERNAL_ASSERT(false, "unsupported expr_type: ", static_cast(expr_type)); return nullptr; } } diff --git a/torch/csrc/jit/tensorexpr/registerizer.cpp b/torch/csrc/jit/tensorexpr/registerizer.cpp index 9ad44e31a3873..acf216c51b530 100644 --- a/torch/csrc/jit/tensorexpr/registerizer.cpp +++ b/torch/csrc/jit/tensorexpr/registerizer.cpp @@ -161,17 +161,12 @@ AccessHashMap& Scope::getAccessMapByBuf(const BufPtr& b) { } void Scope::filterClosed() { - closedAccesses_.erase( - std::remove_if( - closedAccesses_.begin(), - closedAccesses_.end(), - [](auto info) { - return info->store_cost()->isConstant() && - immediateAs(info->store_cost()) <= 1 && - info->load_cost()->isConstant() && - immediateAs(info->load_cost()) <= 1; - }), - closedAccesses_.end()); + std::erase_if(closedAccesses_, [](auto info) { + return info->store_cost()->isConstant() && + immediateAs(info->store_cost()) <= 1 && + info->load_cost()->isConstant() && + immediateAs(info->load_cost()) <= 1; + }); } // RegisterizerAnalysis diff --git a/torch/csrc/profiler/collection.cpp b/torch/csrc/profiler/collection.cpp index 617316617fc67..1be2d80310910 100644 --- a/torch/csrc/profiler/collection.cpp +++ b/torch/csrc/profiler/collection.cpp @@ -10,6 +10,7 @@ #include #include +#include #ifdef USE_KINETO #include @@ -1096,6 +1097,18 @@ class TransferEvents { for (const auto* activity : trace_activities_) { auto e = toResult(activity); if (e) { + // Flow data for Kineto events is already set during + // resultFromActivity(). TorchOp events need it copied here because + // their Result is created during RecordFunction callbacks, before + // flow data exists on the GenericTraceActivity. + e->visit(c10::overloaded( + [&](ExtraFields& i) { + i.flow = { + /*id=*/static_cast(activity->flowId()), + /*type=*/static_cast(activity->flowType()), + /*start=*/activity->flowStart()}; + }, + [](auto&) {})); if (config_.experimental_config.expose_kineto_event_metadata) { e->visit(c10::overloaded( [&](ExtraFields& i) { @@ -1105,6 +1118,27 @@ class TransferEvents { i.metadata_json_ = activity->metadataJson(); }, [](auto&) { return; })); + // Parse metadataJson() into extra_meta_ so events() exposes + // Kineto metadata as typed fields without export_chrome_trace(). + // Python schemas (profiler_util.py) are the single SOT for + // which keys to expose and how to type-convert them. + e->visit(c10::overloaded( + [&](ExtraFields& i) { + auto json_str = activity->metadataJson(); + if (!json_str.empty()) { + auto j = nlohmann::json::parse( + "{" + json_str + "}", nullptr, false); + if (!j.is_discarded()) { + for (auto& [key, val] : j.items()) { + i.extra_meta_.emplace( + key, + val.is_string() ? val.get() + : val.dump()); + } + } + } + }, + [](auto&) {})); } const auto* linked_activity = activity->linkedActivity(); if (linked_activity) { diff --git a/torch/csrc/profiler/collection.h b/torch/csrc/profiler/collection.h index d66eb630a47d9..bd46814f8be9b 100644 --- a/torch/csrc/profiler/collection.h +++ b/torch/csrc/profiler/collection.h @@ -132,6 +132,14 @@ using extra_args_t = std::unordered_map; using extra_meta_t = std::unordered_map; using kwinputs_t = std::unordered_map; +// Mirrors `libkineto::GenericTraceActivity::Flow`. Used during post processing +// to embed Kineto events into the broader profiler tree structure. +struct Flow { + uint32_t id{0}; + uint32_t type{0}; + uint32_t start{0}; +}; + struct FallbackPair { ProfilerVoidEventStub device_event_start_ = nullptr; ProfilerVoidEventStub device_event_end_ = nullptr; @@ -179,6 +187,7 @@ struct ExtraFields : TorchOpBasicFields { bool allow_tf32_cublas_; std::unique_ptr perf_event_counters_; std::string metadata_json_; + Flow flow; }; template <> @@ -354,16 +363,6 @@ struct ExtraFields : public PyExtraFieldsBase { template <> struct ExtraFields { - // Mirrors `libkineto::GenericTraceActivity::Flow`. This information is used - // during post processing to properly embed Kineto events into the broader - // profiler tree structure. End users are not generally expected to use these - // fields directly, but they are available for debugging. - struct Flow { - uint32_t id{0}; - uint32_t type{0}; - uint32_t start{0}; - }; - std::string name_; int64_t duration_ns_{0}; uint64_t correlation_id_{0}; @@ -371,6 +370,7 @@ struct ExtraFields { Flow flow; std::weak_ptr linked_activity_; std::string metadata_json_; + extra_meta_t extra_meta_; }; struct TORCH_API Result : public std::enable_shared_from_this { diff --git a/torch/csrc/profiler/data_flow.cpp b/torch/csrc/profiler/data_flow.cpp index a9f98930f8c66..e18d88960fbba 100644 --- a/torch/csrc/profiler/data_flow.cpp +++ b/torch/csrc/profiler/data_flow.cpp @@ -131,15 +131,10 @@ void calculateUniqueTensorIDs( tensor_set.insert(t.allocation_id_ref_.get().value()); } } - tensors.erase( - std::remove_if( - tensors.begin(), - tensors.end(), - [&tensor_set](const auto& i) { - auto it = tensor_set.find(i.allocation_id_ref_.get().value()); - return it == tensor_set.end(); - }), - tensors.end()); + std::erase_if(tensors, [&tensor_set](const auto& i) { + auto it = tensor_set.find(i.allocation_id_ref_.get().value()); + return it == tensor_set.end(); + }); } // Handle the case that the storage of a TensorImpl changed. diff --git a/torch/csrc/profiler/unwind/action.h b/torch/csrc/profiler/unwind/action.h index 5a982cfd046a0..7dc6c6c0b8166 100644 --- a/torch/csrc/profiler/unwind/action.h +++ b/torch/csrc/profiler/unwind/action.h @@ -11,7 +11,8 @@ enum { A_REG_PLUS_DATA_DEREF = 0x3 // exp = *(REG[reg] + data0) }; -// register numbers in dwarf info +// DWARF register numbers — architecture-specific +#if defined(__x86_64__) enum { D_UNDEFINED = -1, D_RBP = 6, @@ -19,6 +20,28 @@ enum { D_RIP = 16, D_REG_SIZE = 17, }; +static constexpr int D_FRAME_PTR = D_RBP; +static constexpr int D_STACK_PTR = D_RSP; +static constexpr int D_RET_ADDR = D_RIP; +static constexpr int D_EXPECTED_RA_REG = 16; +#elif defined(__aarch64__) +enum { + D_UNDEFINED = -1, + D_FP = 29, + D_LR = 30, + D_SP = 31, + D_REG_SIZE = 32, +}; +static constexpr int D_FRAME_PTR = D_FP; +static constexpr int D_STACK_PTR = D_SP; +static constexpr int D_RET_ADDR = D_LR; +static constexpr int D_EXPECTED_RA_REG = 30; +#else +enum { + D_UNDEFINED = -1, + D_REG_SIZE = 1, +}; +#endif struct Action { uint8_t kind = A_UNDEFINED; diff --git a/torch/csrc/profiler/unwind/dwarf_enums.h b/torch/csrc/profiler/unwind/dwarf_enums.h index 91af24b34e1f9..3effefd1f0083 100644 --- a/torch/csrc/profiler/unwind/dwarf_enums.h +++ b/torch/csrc/profiler/unwind/dwarf_enums.h @@ -30,6 +30,7 @@ enum { DW_CFA_advance_loc1 = 0x02, DW_CFA_advance_loc2 = 0x03, DW_CFA_advance_loc4 = 0x04, + DW_CFA_offset_extended = 0x05, DW_CFA_restore_extended = 0x06, DW_CFA_undefined = 0x07, DW_CFA_register = 0x09, diff --git a/torch/csrc/profiler/unwind/fde.h b/torch/csrc/profiler/unwind/fde.h index ffb06b5ab1f46..70b829a0f8023 100644 --- a/torch/csrc/profiler/unwind/fde.h +++ b/torch/csrc/profiler/unwind/fde.h @@ -63,8 +63,10 @@ struct FDE { } else { ra_register_ = static_cast(LC.readULEB128()); } - // we assume this in the state - TORCH_INTERNAL_ASSERT(ra_register_ == 16, "unexpected number of registers"); + TORCH_INTERNAL_ASSERT( + ra_register_ == D_EXPECTED_RA_REG, + "unexpected ra register: ", + ra_register_); if (augmentation_string_ && *augmentation_string_ == 'z') { augmentation_length_ = static_cast(LC.readULEB128()); Lexer A(LC.loc()); @@ -271,6 +273,11 @@ struct FDE { auto delta = L.read(); return advance_loc(delta); } + case DW_CFA_offset_extended: { + auto reg = L.readULEB128(); + auto off = L.readULEB128(); + return offset(reg, off); + } case DW_CFA_restore_extended: { auto reg = L.readULEB128(); return restore(reg); diff --git a/torch/csrc/profiler/unwind/unwind.cpp b/torch/csrc/profiler/unwind/unwind.cpp index bdb610627c897..f779af0d1ade7 100644 --- a/torch/csrc/profiler/unwind/unwind.cpp +++ b/torch/csrc/profiler/unwind/unwind.cpp @@ -3,32 +3,32 @@ #include #include -#if !defined(__linux__) || !defined(__x86_64__) || !defined(__has_include) || \ - !__has_include("ext/stdio_filebuf.h") +#if !defined(__linux__) || !(defined(__x86_64__) || defined(__aarch64__)) || \ + !defined(__has_include) || !__has_include("ext/stdio_filebuf.h") namespace torch::unwind { std::vector unwind() { TORCH_WARN_ONCE( - "record_context_cpp is not support on non-linux non-x86_64 platforms"); + "record_context_cpp is not supported on this platform (requires linux x86_64 or aarch64)"); return {}; } std::optional> libraryFor(void* addr) { TORCH_WARN_ONCE( - "record_context_cpp is not support on non-linux non-x86_64 platforms"); + "record_context_cpp is not supported on this platform (requires linux x86_64 or aarch64)"); return {}; } #ifndef FBCODE_CAFFE2 std::vector symbolize(const std::vector& frames, Mode mode) { TORCH_WARN_ONCE( - "record_context_cpp is not support on non-linux non-x86_64 platforms"); + "record_context_cpp is not supported on this platform (requires linux x86_64 or aarch64)"); return {}; } #endif Stats stats() { TORCH_WARN_ONCE( - "record_context_cpp is not support on non-linux non-x86_64 platforms"); + "record_context_cpp is not supported on this platform (requires linux x86_64 or aarch64)"); return {}; } @@ -41,8 +41,10 @@ Stats stats() { #include #include #include +#include #include #include +#include #include #include @@ -53,7 +55,14 @@ Stats stats() { #include #include +#if defined(__aarch64__) +extern "C" void unwind_c( + std::vector* result, + uintptr_t fp, + uintptr_t lr); +#else extern "C" void unwind_c(std::vector* result, int64_t rsp, int64_t rbp); +#endif extern "C" void unwind_entry(std::vector* result); namespace torch::unwind { @@ -98,7 +107,8 @@ struct LibraryInfo { void* fde_data = eh_frame_hdr_.entryForAddr(addr); FDE fde(fde_data, name().c_str(), load_bias()); TableState state = fde.readUpTo(addr); - return Unwinder(state.cfa, state.registers[D_RIP], state.registers[D_RBP]); + return Unwinder( + state.cfa, state.registers[D_RET_ADDR], state.registers[D_FRAME_PTR]); } const std::string& name() const { return name_; @@ -501,6 +511,82 @@ Stats stats() { } // namespace torch::unwind +#if defined(__aarch64__) +// aarch64 uses frame-pointer chain walking instead of DWARF unwinding. +// Each frame has: *(FP) = caller's FP, *(FP+8) = saved LR (return address). +// This is simpler and avoids issues with tail calls producing stale x30 +// values in DWARF-based unwinding. GCC/Clang on aarch64 emit frame +// pointers by default even at -O2. +// +// External libraries (CPython, libc, CUDA runtime) may be built without +// frame pointers, making x29 an arbitrary callee-saved value. We obtain +// the current thread's stack bounds and reject any fp outside that range +// to avoid dereferencing garbage pointers. +// +// No cache_mutex_ needed: frame-pointer walking reads only the stack, +// unlike the x86 path which queries the DWARF FDE cache. + +static bool get_stack_bounds(uintptr_t& lo, uintptr_t& hi) { + pthread_attr_t attr; + if (pthread_getattr_np(pthread_self(), &attr) != 0) { + return false; + } + void* base = nullptr; + size_t size = 0; + int rc = pthread_attr_getstack(&attr, &base, &size); + pthread_attr_destroy(&attr); + if (rc != 0) { + return false; + } + lo = reinterpret_cast(base); + hi = lo + size; + return true; +} + +extern "C" C10_USED void unwind_c( + std::vector* result, + uintptr_t fp, + uintptr_t lr) { + // NOLINTNEXTLINE(performance-no-int-to-ptr) + result->push_back((void*)lr); + + uintptr_t stack_lo = 0, stack_hi = 0; + if (!get_stack_bounds(stack_lo, stack_hi)) { + return; + } + + constexpr int kMaxFrames = 4096; + int depth = 0; + while (fp != 0 && (fp & 0xF) == 0 && depth++ < kMaxFrames) { + if (fp < stack_lo || fp + 16 > stack_hi) { + break; + } + uintptr_t saved_lr; + std::memcpy( + &saved_lr, reinterpret_cast(fp + 8), sizeof(saved_lr)); + if (saved_lr == 0) { + break; + } + // NOLINTNEXTLINE(performance-no-int-to-ptr) + result->push_back((void*)saved_lr); + uintptr_t next_fp; + std::memcpy(&next_fp, reinterpret_cast(fp), sizeof(next_fp)); + if (next_fp <= fp) { + break; + } + fp = next_fp; + } +} + +// x0 already holds the result pointer. +// Pass FP (x29) and LR (x30), then tail-call unwind_c. +__asm__( + ".global unwind_entry\n" + "unwind_entry:\n" + "mov x1, x29\n" + "mov x2, x30\n" + "b unwind_c\n"); +#else extern "C" C10_USED void unwind_c( std::vector* result, int64_t rsp, @@ -508,17 +594,17 @@ extern "C" C10_USED void unwind_c( std::shared_lock lock(torch::unwind::cache_mutex_); torch::unwind::UnwindState state{}; // NOLINTNEXTLINE(performance-no-int-to-ptr) - state.rip = *(int64_t*)rsp; + state.pc = *(int64_t*)rsp; // +8 because we saved rsp after the return address was already pushed // to the stack - state.rsp = rsp + 8; - state.rbp = rbp; + state.sp = rsp + 8; + state.fp = rbp; torch::unwind::unwind_cache.checkRefresh(lock); - while (true) { // unwind for _start sets rip as being undefined + while (true) { // NOLINTNEXTLINE(performance-no-int-to-ptr) - result->push_back((void*)state.rip); + result->push_back((void*)state.pc); const torch::unwind::Unwinder& uw = - torch::unwind::unwind_cache.unwinderFor(state.rip, lock); + torch::unwind::unwind_cache.unwinderFor(state.pc, lock); if (uw.terminator()) { if (uw.isUnknown()) { result->push_back(nullptr); @@ -529,8 +615,7 @@ extern "C" C10_USED void unwind_c( } } -// calling convention puts the first three pointer/int64_t arguments in -// rdi rsi rdx (all caller-saved) +// x86-64 calling convention: rdi rsi rdx (all caller-saved) // rdi already holds the pointer to the result vector // we add arguments for current rsp and rbp and then tail call // into unwind_c @@ -540,5 +625,6 @@ __asm__( "mov %rsp, %rsi;\n" "mov %rbp, %rdx;\n" "jmp unwind_c;\n"); +#endif #endif diff --git a/torch/csrc/profiler/unwind/unwinder.h b/torch/csrc/profiler/unwind/unwinder.h index d673f47af8db2..3836203688530 100644 --- a/torch/csrc/profiler/unwind/unwinder.h +++ b/torch/csrc/profiler/unwind/unwinder.h @@ -6,30 +6,32 @@ namespace torch::unwind { +// Architecture-neutral names: pc (program counter / return address), +// fp (frame pointer: x86 RBP, aarch64 x29), sp (stack pointer). struct UnwindState { - int64_t rip, rbp, rsp; + int64_t pc, fp, sp; }; struct Unwinder { - Unwinder(Action rsp, Action rip, Action rbp) - : kind_(rip.kind == A_UNDEFINED ? END : STANDARD), - reg_(rsp.reg), - off_(rsp.data), - rip_off_(rip.data), - rbp_off_( - rbp.kind == A_UNDEFINED ? std::numeric_limits::max() - : rbp.data), - deref_(rsp.kind == A_REG_PLUS_DATA_DEREF) { - check(rsp.reg == D_RSP || rsp.reg == D_RBP); - check(rip.kind == A_UNDEFINED || rip.kind == A_LOAD_CFA_OFFSET); - if (rsp.kind == A_REG_PLUS_DATA) { - check(rbp.kind == A_LOAD_CFA_OFFSET || rbp.kind == A_UNDEFINED); - } else if (rsp.kind == A_REG_PLUS_DATA_DEREF) { - if (rbp.kind == A_REG_PLUS_DATA_DEREF) { - check(rbp.reg == rsp.reg); - rbp_off_ -= rsp.data; + Unwinder(Action cfa, Action ret, Action fp) + : kind_(ret.kind == A_UNDEFINED ? END : STANDARD), + reg_(cfa.reg), + off_(cfa.data), + ret_off_(ret.data), + fp_off_( + fp.kind == A_UNDEFINED ? std::numeric_limits::max() + : fp.data), + deref_(cfa.kind == A_REG_PLUS_DATA_DEREF) { + check(cfa.reg == D_STACK_PTR || cfa.reg == D_FRAME_PTR); + check(ret.kind == A_UNDEFINED || ret.kind == A_LOAD_CFA_OFFSET); + if (cfa.kind == A_REG_PLUS_DATA) { + check(fp.kind == A_LOAD_CFA_OFFSET || fp.kind == A_UNDEFINED); + } else if (cfa.kind == A_REG_PLUS_DATA_DEREF) { + if (fp.kind == A_REG_PLUS_DATA_DEREF) { + check(fp.reg == cfa.reg); + fp_off_ -= cfa.data; } else { - check(rbp.kind == A_UNDEFINED); + check(fp.kind == A_UNDEFINED); } } else { check(false); @@ -53,28 +55,28 @@ struct Unwinder { } UnwindState run(const UnwindState& cur) const { UnwindState r = cur; - r.rsp = (reg_ == D_RSP ? cur.rsp : cur.rbp) + off_; - r.rbp = rbp_off_ == std::numeric_limits::max() - ? cur.rbp + r.sp = (reg_ == D_STACK_PTR ? cur.sp : cur.fp) + off_; + r.fp = fp_off_ == std::numeric_limits::max() + ? cur.fp // NOLINTNEXTLINE(performance-no-int-to-ptr) - : *(int64_t*)(r.rsp + rbp_off_); + : *(int64_t*)(r.sp + fp_off_); if (deref_) { // NOLINTNEXTLINE(performance-no-int-to-ptr) - r.rsp = *(int64_t*)r.rsp; + r.sp = *(int64_t*)r.sp; } // NOLINTNEXTLINE(performance-no-int-to-ptr) - r.rip = *(int64_t*)(r.rsp + rip_off_); + r.pc = *(int64_t*)(r.sp + ret_off_); return r; } private: - Unwinder() : kind_(UNKNOWN), reg_(0), off_(0), rip_off_(0), rbp_off_(0) {} + Unwinder() : kind_(UNKNOWN), reg_(0), off_(0), ret_off_(0), fp_off_(0) {} enum Kind { STANDARD, END, UNKNOWN } kind_; uint32_t reg_; int64_t off_; - int64_t rip_off_; - int64_t rbp_off_; + int64_t ret_off_; + int64_t fp_off_; bool deref_{false}; }; diff --git a/torch/csrc/utils/device_lazy_init.cpp b/torch/csrc/utils/device_lazy_init.cpp index b1338fc43e913..fc5c7a47b8b82 100644 --- a/torch/csrc/utils/device_lazy_init.cpp +++ b/torch/csrc/utils/device_lazy_init.cpp @@ -12,24 +12,21 @@ namespace torch::utils { namespace { -std::array is_initialized{}; -std::array is_in_bad_fork{}; -std::array at_fork_registered{}; +std::array, at::COMPILE_TIME_MAX_DEVICE_TYPES> + is_initialized{}; +std::array, at::COMPILE_TIME_MAX_DEVICE_TYPES> + is_in_bad_fork{}; +std::array, at::COMPILE_TIME_MAX_DEVICE_TYPES> + at_fork_registered{}; c10::once_flag at_fork_register_once{}; } // anonymous namespace bool is_device_initialized(at::DeviceType device_type) { - pybind11::gil_scoped_acquire g; return is_initialized[static_cast(device_type)]; } void device_lazy_init(at::DeviceType device_type) { - pybind11::gil_scoped_acquire g; - // Protected by the GIL. We don't use call_once because under ASAN it - // has a buggy implementation that deadlocks if an instance throws an - // exception. In any case, call_once isn't necessary, because we - // have taken a lock. if (is_device_initialized(device_type)) { return; } @@ -40,6 +37,10 @@ void device_lazy_init(at::DeviceType device_type) { return; } + // Don't use call_once because under ASAN it has a buggy implementation that + // deadlocks if an instance throws an exception and Python _lazy_init() + // functions are idempotent. + pybind11::gil_scoped_acquire g; std::string module_name = "torch." + at::DeviceTypeName(device_type, true); auto module = THPObjectPtr(PyImport_ImportModule(module_name.c_str())); if (!module) { @@ -50,6 +51,7 @@ void device_lazy_init(at::DeviceType device_type) { auto has_lazy_init_method = PyObject_HasAttrString(module.get(), "_lazy_init") == 1; if (!has_lazy_init_method) { + is_initialized[static_cast(device_type)] = true; return; } } diff --git a/torch/csrc/utils/generated_serialization_types.h b/torch/csrc/utils/generated_serialization_types.h index f4614ac0c191a..995e2405cb4da 100644 --- a/torch/csrc/utils/generated_serialization_types.h +++ b/torch/csrc/utils/generated_serialization_types.h @@ -1,5 +1,5 @@ // @generated by update_schema.py -// checksum<> +// checksum<> // clang-format off #pragma once @@ -289,6 +289,8 @@ enum class ScalarType { FLOAT8E4M3FNUZ = 31, FLOAT8E5M2FNUZ = 32, FLOAT8E8M0FNU = 33, + UINT32 = 34, + UINT64 = 35, }; inline std::string_view printEnum(const ScalarType& e) { @@ -313,6 +315,8 @@ inline std::string_view printEnum(const ScalarType& e) { case ScalarType::FLOAT8E4M3FNUZ: return "FLOAT8E4M3FNUZ"; case ScalarType::FLOAT8E5M2FNUZ: return "FLOAT8E5M2FNUZ"; case ScalarType::FLOAT8E8M0FNU: return "FLOAT8E8M0FNU"; + case ScalarType::UINT32: return "UINT32"; + case ScalarType::UINT64: return "UINT64"; default: throw std::runtime_error("Unknown enum value"); } @@ -339,6 +343,8 @@ inline void parseEnum(std::string_view s, ScalarType& t) { if (s == "FLOAT8E4M3FNUZ") { t = ScalarType::FLOAT8E4M3FNUZ; return; } if (s == "FLOAT8E5M2FNUZ") { t = ScalarType::FLOAT8E5M2FNUZ; return; } if (s == "FLOAT8E8M0FNU") { t = ScalarType::FLOAT8E8M0FNU; return; } + if (s == "UINT32") { t = ScalarType::UINT32; return; } + if (s == "UINT64") { t = ScalarType::UINT64; return; } throw std::runtime_error("Unknown enum value: " + std::string{s}); } @@ -1235,11 +1241,11 @@ class Argument { public: enum class Tag { - AS_NONE, AS_TENSOR, AS_TENSORS, AS_INT, AS_INTS, AS_FLOAT, AS_FLOATS, AS_STRING, AS_STRINGS, AS_SYM_INT, AS_SYM_INTS, AS_SCALAR_TYPE, AS_MEMORY_FORMAT, AS_LAYOUT, AS_DEVICE, AS_BOOL, AS_BOOLS, AS_SYM_BOOL, AS_SYM_BOOLS, AS_GRAPH, AS_OPTIONAL_TENSORS, AS_CUSTOM_OBJ, AS_OPERATOR, AS_SYM_FLOAT, AS_SYM_FLOATS, AS_OPTIONAL_TENSOR, AS_COMPLEX, AS_NESTED_TENSORS, AS_INT_LISTS, AS_STRING_TO_ARGUMENT + AS_NONE, AS_TENSOR, AS_TENSORS, AS_INT, AS_INTS, AS_FLOAT, AS_FLOATS, AS_STRING, AS_STRINGS, AS_SYM_INT, AS_SYM_INTS, AS_SCALAR_TYPE, AS_MEMORY_FORMAT, AS_LAYOUT, AS_DEVICE, AS_BOOL, AS_BOOLS, AS_SYM_BOOL, AS_SYM_BOOLS, AS_GRAPH, AS_OPTIONAL_TENSORS, AS_CUSTOM_OBJ, AS_OPERATOR, AS_SYM_FLOAT, AS_SYM_FLOATS, AS_OPTIONAL_TENSOR, AS_COMPLEX, AS_NESTED_TENSORS, AS_INT_LISTS, AS_STRING_TO_ARGUMENT, AS_FLOAT_LISTS }; private: - std::variant, int64_t, std::vector, F64, std::vector, std::string, std::vector, SymIntArgument, std::vector, ScalarType, MemoryFormat, Layout, Device, bool, std::vector, SymBoolArgument, std::vector, GraphArgument, std::vector, CustomObjArgument, std::string, SymFloatArgument, std::vector, OptionalTensorArgument, ComplexValue, std::vector>, std::vector>, std::unordered_map>> variant_; + std::variant, int64_t, std::vector, F64, std::vector, std::string, std::vector, SymIntArgument, std::vector, ScalarType, MemoryFormat, Layout, Device, bool, std::vector, SymBoolArgument, std::vector, GraphArgument, std::vector, CustomObjArgument, std::string, SymFloatArgument, std::vector, OptionalTensorArgument, ComplexValue, std::vector>, std::vector>, std::unordered_map>, std::vector>> variant_; Tag tag_; public: @@ -1517,6 +1523,15 @@ class Argument { tag_ = Tag::AS_STRING_TO_ARGUMENT; } + const std::vector>& get_as_float_lists() const { + return std::get<31>(variant_); + } + + void set_as_float_lists(std::vector> def) { + variant_.emplace<31>(std::move(def)); + tag_ = Tag::AS_FLOAT_LISTS; + } + friend void to_json(nlohmann::json& nlohmann_json_j, const Argument& nlohmann_json_t) { if (nlohmann_json_t.tag_ == Tag::AS_NONE) { @@ -1639,6 +1654,10 @@ class Argument { nlohmann_json_j["as_string_to_argument"] = nlohmann_json_t.get_as_string_to_argument(); return; } + if (nlohmann_json_t.tag_ == Tag::AS_FLOAT_LISTS) { + nlohmann_json_j["as_float_lists"] = nlohmann_json_t.get_as_float_lists(); + return; + } } friend void from_json(const nlohmann::json& nlohmann_json_j, Argument& nlohmann_json_t) { @@ -1793,6 +1812,11 @@ class Argument { nlohmann_json_t.tag_ = Tag::AS_STRING_TO_ARGUMENT; return; } + if (nlohmann_json_j.contains("as_float_lists")) { + nlohmann_json_t.variant_.emplace<31>(nlohmann_json_j.at("as_float_lists").template get>>()); + nlohmann_json_t.tag_ = Tag::AS_FLOAT_LISTS; + return; + } } }; @@ -1828,6 +1852,7 @@ inline std::string_view printEnum(const Argument::Tag& e) { case Argument::Tag::AS_NESTED_TENSORS: return "AS_NESTED_TENSORS"; case Argument::Tag::AS_INT_LISTS: return "AS_INT_LISTS"; case Argument::Tag::AS_STRING_TO_ARGUMENT: return "AS_STRING_TO_ARGUMENT"; + case Argument::Tag::AS_FLOAT_LISTS: return "AS_FLOAT_LISTS"; default: throw std::runtime_error("Unknown enum value"); } @@ -1864,6 +1889,7 @@ inline void parseEnum(std::string_view s, Argument::Tag& t) { if (s == "AS_NESTED_TENSORS") { t = Argument::Tag::AS_NESTED_TENSORS; return; } if (s == "AS_INT_LISTS") { t = Argument::Tag::AS_INT_LISTS; return; } if (s == "AS_STRING_TO_ARGUMENT") { t = Argument::Tag::AS_STRING_TO_ARGUMENT; return; } + if (s == "AS_FLOAT_LISTS") { t = Argument::Tag::AS_FLOAT_LISTS; return; } throw std::runtime_error("Unknown enum value: " + std::string{s}); } diff --git a/torch/csrc/utils/python_arg_parser.cpp b/torch/csrc/utils/python_arg_parser.cpp index c512a91833aa0..a9f52cafb65ad 100644 --- a/torch/csrc/utils/python_arg_parser.cpp +++ b/torch/csrc/utils/python_arg_parser.cpp @@ -769,9 +769,12 @@ auto handle_torch_function_indexing( } py::object func = PyObject_FastGetAttrString(THPVariableClass, (char*)func_name); - py::object args = (val == nullptr) - ? py::make_tuple(py::handle(self), py::handle(index)) - : py::make_tuple(py::handle(self), py::handle(index), py::handle(val)); + py::tuple args; + if (val == nullptr) { + args = py::make_tuple(py::handle(self), py::handle(index)); + } else { + args = py::make_tuple(py::handle(self), py::handle(index), py::handle(val)); + } return handle_torch_function_no_python_arg_parser( overridable_args, args.ptr(), @@ -892,7 +895,7 @@ bool is_tensor_and_append_overloaded( static bool is_scalar_list( PyObject* obj, std::vector* overloaded_args = nullptr) { - auto tuple = six::isTuple(obj); + auto tuple = PyTuple_Check(obj); if (!(tuple || PyList_Check(obj))) { return false; } @@ -923,7 +926,7 @@ bool is_tensor_list_and_append_overloaded( std::vector* overloaded_args, size_t argnum, bool throw_error) { - auto tuple = six::isTuple(obj); + auto tuple = PyTuple_Check(obj); if (!(tuple || PyList_Check(obj))) { return false; } @@ -964,7 +967,7 @@ static bool is_float_or_symfloat(PyObject* obj) { static bool is_float_or_complex_list( PyObject* obj, std::vector* overloaded_args = nullptr) { - auto tuple = six::isTuple(obj); + auto tuple = PyTuple_Check(obj); if (!(tuple || PyList_Check(obj))) { return false; } diff --git a/torch/csrc/utils/python_arg_parser.h b/torch/csrc/utils/python_arg_parser.h index 12b023c764749..c72d3046a90dc 100644 --- a/torch/csrc/utils/python_arg_parser.h +++ b/torch/csrc/utils/python_arg_parser.h @@ -65,7 +65,7 @@ #include #include #include -#include +#include #include #include @@ -420,14 +420,14 @@ inline at::Scalar PythonArgs::scalar(int i) { inline std::vector PythonArgs::scalarlist(int i) { if (!args[i]) return std::vector(); - auto tuple = six::isTuple(args[i]); - THPObjectPtr arg = six::maybeAsTuple(args[i]); + auto tuple = PyTuple_Check(args[i]); + TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i])); // NOLINTNEXTLINE(bugprone-branch-clone) - auto size = tuple ? PyTuple_GET_SIZE(arg.get()) : PyList_GET_SIZE(arg.get()); + auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]); std::vector res(size); for (const auto idx : c10::irange(size)) { - PyObject* obj = tuple ? PyTuple_GET_ITEM(arg.get(), idx) - : PyList_GET_ITEM(arg.get(), idx); + PyObject* obj = + tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx); res[idx] = scalar_slow(obj); } return res; @@ -450,14 +450,14 @@ inline std::optional PythonArgs::scalarOptional(int i) { inline std::vector PythonArgs::tensorlist(int i) { if (!args[i]) return std::vector(); - auto tuple = six::isTuple(args[i]); - THPObjectPtr arg = six::maybeAsTuple(args[i]); + auto tuple = PyTuple_Check(args[i]); + TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i])); // NOLINTNEXTLINE(bugprone-branch-clone) - auto size = tuple ? PyTuple_GET_SIZE(arg.get()) : PyList_GET_SIZE(arg.get()); + auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]); std::vector res(size); for (const auto idx : c10::irange(size)) { - PyObject* obj = tuple ? PyTuple_GET_ITEM(arg.get(), idx) - : PyList_GET_ITEM(arg.get(), idx); + PyObject* obj = + tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx); // This is checked by the argument parser so it's safe to cast without // checking if this is a tensor first res[idx] = THPVariable_Unpack(obj); @@ -469,15 +469,15 @@ inline torch::List> PythonArgs:: list_of_optional_tensors(int i) { if (!args[i]) return torch::List>(); - auto tuple = six::isTuple(args[i]); - THPObjectPtr arg = six::maybeAsTuple(args[i]); + auto tuple = PyTuple_Check(args[i]); + TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i])); // NOLINTNEXTLINE(bugprone-branch-clone) - auto size = tuple ? PyTuple_GET_SIZE(arg.get()) : PyList_GET_SIZE(arg.get()); + auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]); torch::List> res; res.reserve(size); for (const auto idx : c10::irange(size)) { - PyObject* obj = tuple ? PyTuple_GET_ITEM(arg.get(), idx) - : PyList_GET_ITEM(arg.get(), idx); + PyObject* obj = + tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx); // This is checked by the argument parser so it's safe to cast without // checking if this is a tensor first res.push_back(THPVariable_Unpack(obj)); @@ -490,18 +490,18 @@ inline std::array PythonArgs::tensorlist_n(int i) { auto res = std::array(); if (!args[i]) return res; - auto tuple = six::isTuple(args[i]); - THPObjectPtr arg = six::maybeAsTuple(args[i]); + auto tuple = PyTuple_Check(args[i]); + TORCH_INTERNAL_ASSERT_DEBUG_ONLY(tuple || PyList_Check(args[i])); // NOLINTNEXTLINE(bugprone-branch-clone) - auto size = tuple ? PyTuple_GET_SIZE(arg.get()) : PyList_GET_SIZE(arg.get()); + auto size = tuple ? PyTuple_GET_SIZE(args[i]) : PyList_GET_SIZE(args[i]); if (size != N) { TORCH_CHECK_TYPE( false, fmt::format("expected tuple of {} elements but got {}", N, size)); } for (const auto idx : c10::irange(size)) { - PyObject* obj = tuple ? PyTuple_GET_ITEM(arg.get(), idx) - : PyList_GET_ITEM(arg.get(), idx); + PyObject* obj = + tuple ? PyTuple_GET_ITEM(args[i], idx) : PyList_GET_ITEM(args[i], idx); // This is checked by the argument parser so it's safe to cast without // checking if this is a tensor first res[idx] = THPVariable_Unpack(obj); diff --git a/torch/csrc/utils/python_dispatch.cpp b/torch/csrc/utils/python_dispatch.cpp index 0be7c4a3970af..c141aa928307a 100644 --- a/torch/csrc/utils/python_dispatch.cpp +++ b/torch/csrc/utils/python_dispatch.cpp @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -30,15 +31,18 @@ namespace py = pybind11; namespace torch::impl::dispatch { // Global storage for leaked Python filenames to ensure they remain valid -// for the lifetime of Library objects -static std::vector leaked_python_filenames_; +// for the lifetime of Library objects. We use unique_ptr rather than +// plain string so that c_str() pointers handed to Library objects remain valid +// when the vector reallocates. +static c10::Synchronized>> + leaked_python_filenames_; // NB: I'd like to index this on OperatorHandle, but I can't, as I can't // guarantee that the main interpreter has finish doing all registrations before // the other interpreters start banging on it -static ska::flat_hash_map< +static c10::Synchronized>> + ska::flat_hash_map>>> python_registrations_; static torch::Library::Kind parseKind(const std::string& k) { @@ -424,10 +428,12 @@ void initDispatchBindings(PyObject* module) { std::make_unique( func, dispatch, with_keyset))), register_or_verify()); - python_registrations_[lib._resolve(name)].insert_or_assign( - dispatch, - std::make_shared( - func.release().ptr(), getPyInterpreter())); + python_registrations_.withLock([&](auto& regs) { + regs[lib._resolve(name)].insert_or_assign( + dispatch, + std::make_shared( + func.release().ptr(), getPyInterpreter())); + }); } END_HANDLE_TH_ERRORS_PYBIND }, @@ -512,8 +518,11 @@ void initDispatchBindings(PyObject* module) { HANDLE_TH_ERRORS // Store the file string in global storage to ensure it remains valid // for the lifetime of the Library object - leaked_python_filenames_.emplace_back(file); - const char* leaked_file = leaked_python_filenames_.back().c_str(); + const char* leaked_file = + leaked_python_filenames_.withLock([&](auto& filenames) { + filenames.push_back(std::make_unique(file)); + return filenames.back()->c_str(); + }); return std::make_unique( parseKind(kind), @@ -534,7 +543,7 @@ void initDispatchBindings(PyObject* module) { m.def( "_dispatch_clear_leaked_python_filenames", - []() { leaked_python_filenames_.clear(); }, + []() { leaked_python_filenames_.withLock([](auto& f) { f.clear(); }); }, "Clear the global storage of leaked Python filenames. " "WARNING: Only call this if you're sure no Library objects are still using the filenames."); @@ -1077,7 +1086,8 @@ void python_op_registration_trampoline_impl( auto arguments = torch::jit::pop(*stack, op.schema().arguments().size()); py::gil_scoped_acquire g; auto args_kwargs = parseIValuesToPyArgsKwargs(op, arguments); - const auto& func = python_registrations_[op.operator_name()][key]; + auto func = python_registrations_.withLock( + [&](auto& regs) { return regs[op.operator_name()][key]; }); TORCH_INTERNAL_ASSERT(func != nullptr); auto* pyobj = func->ptr(getPyInterpreter()); TORCH_INTERNAL_ASSERT(pyobj != nullptr); diff --git a/torch/csrc/utils/six.h b/torch/csrc/utils/six.h deleted file mode 100644 index 9671e5156b9d4..0000000000000 --- a/torch/csrc/utils/six.h +++ /dev/null @@ -1,52 +0,0 @@ -#pragma once - -#include -#include -#include -#include - -namespace six { - -// Usually instances of PyStructSequence is also an instance of tuple -// but in some py2 environment it is not, so we have to manually check -// the name of the type to determine if it is a namedtupled returned -// by a pytorch operator. - -inline bool isStructSeq(pybind11::handle input) { - return pybind11::cast(pybind11::type::handle_of(input).attr( - "__module__")) == "torch.return_types"; -} - -inline bool isStructSeq(PyObject* obj) { - return isStructSeq(pybind11::handle(obj)); -} - -inline bool isTuple(pybind11::handle input) { - if (PyTuple_Check(input.ptr())) { - return true; - } - return false; -} - -inline bool isTuple(PyObject* obj) { - return isTuple(pybind11::handle(obj)); -} - -// maybeAsTuple: if the input is a structseq, then convert it to a tuple -// -// On Python 3, structseq is a subtype of tuple, so these APIs could be used -// directly. But on Python 2, structseq is not a subtype of tuple, so we need to -// manually create a new tuple object from structseq. -inline THPObjectPtr maybeAsTuple(PyStructSequence* obj) { - Py_INCREF(obj); - return THPObjectPtr((PyObject*)obj); -} - -inline THPObjectPtr maybeAsTuple(PyObject* obj) { - if (isStructSeq(obj)) - return maybeAsTuple((PyStructSequence*)obj); - Py_INCREF(obj); - return THPObjectPtr(obj); -} - -} // namespace six diff --git a/torch/csrc/utils/structseq.cpp b/torch/csrc/utils/structseq.cpp index 2e804aa44bad9..ded49128cf284 100644 --- a/torch/csrc/utils/structseq.cpp +++ b/torch/csrc/utils/structseq.cpp @@ -12,7 +12,7 @@ * https://github.com/python/cpython#copyright-and-license-information */ -#include +#include #include #include @@ -20,19 +20,12 @@ namespace torch::utils { -// NOTE: The built-in repr method from PyStructSequence was updated in -// https://github.com/python/cpython/commit/c70ab02df2894c34da2223fc3798c0404b41fd79 -// so this function might not be required in Python 3.8+. PyObject* returned_structseq_repr(PyStructSequence* obj) { PyTypeObject* typ = Py_TYPE(obj); - THPObjectPtr tup = six::maybeAsTuple(obj); - if (tup == nullptr) { - return nullptr; - } + Py_ssize_t num_elements = PyTuple_GET_SIZE(obj); std::stringstream ss; ss << typ->tp_name << "(\n"; - Py_ssize_t num_elements = Py_SIZE(obj); for (Py_ssize_t i = 0; i < num_elements; i++) { const char* cname = typ->tp_members[i].name; @@ -46,11 +39,7 @@ PyObject* returned_structseq_repr(PyStructSequence* obj) { return nullptr; } - PyObject* val = PyTuple_GetItem(tup.get(), i); - if (val == nullptr) { - return nullptr; - } - + PyObject* val = PyTuple_GET_ITEM(obj, i); auto repr = THPObjectPtr(PyObject_Repr(val)); if (repr == nullptr) { return nullptr; diff --git a/torch/csrc/utils/tensor_dtypes.cpp b/torch/csrc/utils/tensor_dtypes.cpp index 39df9be68868a..938c7bc7558d0 100644 --- a/torch/csrc/utils/tensor_dtypes.cpp +++ b/torch/csrc/utils/tensor_dtypes.cpp @@ -20,17 +20,15 @@ void initializeDtypes() { for (at::ScalarType scalarType : all_scalar_types) { auto [primary_name, legacy_name] = c10::getDtypeNames(scalarType); - PyObject* dtype = THPDtype_New(scalarType, primary_name); - torch::registerDtypeObject((THPDtype*)dtype, scalarType); - Py_INCREF(dtype); - if (PyModule_AddObject(torch_module.get(), primary_name.c_str(), dtype) != - 0) { + THPObjectPtr dtype(THPDtype_New(scalarType, primary_name)); + torch::registerDtypeObject((THPDtype*)dtype.get(), scalarType); + if (PyModule_AddObjectRef( + torch_module.get(), primary_name.c_str(), dtype.get()) != 0) { throw python_error(); } if (!legacy_name.empty()) { - Py_INCREF(dtype); - if (PyModule_AddObject(torch_module.get(), legacy_name.c_str(), dtype) != - 0) { + if (PyModule_AddObjectRef( + torch_module.get(), legacy_name.c_str(), dtype.get()) != 0) { throw python_error(); } } diff --git a/torch/csrc/utils/tensor_layouts.cpp b/torch/csrc/utils/tensor_layouts.cpp index d0bccbcf9106f..3b1e04cb08a30 100644 --- a/torch/csrc/utils/tensor_layouts.cpp +++ b/torch/csrc/utils/tensor_layouts.cpp @@ -6,49 +6,36 @@ namespace torch::utils { -#define REGISTER_LAYOUT(layout, LAYOUT) \ - PyObject* layout##_layout = \ - THPLayout_New(at::Layout::LAYOUT, "torch." #layout); \ - Py_INCREF(layout##_layout); \ - if (PyModule_AddObject(torch_module, "" #layout, layout##_layout) != 0) { \ - throw python_error(); \ - } \ - registerLayoutObject((THPLayout*)layout##_layout, at::Layout::LAYOUT); +static void registerLayout( + PyObject* torch_module, + at::Layout layout, + const char* name, + const char* qualified_name) { + THPObjectPtr obj(THPLayout_New(layout, qualified_name)); + if (PyModule_AddObjectRef(torch_module, name, obj.get()) != 0) { + throw python_error(); + } + registerLayoutObject((THPLayout*)obj.get(), layout); +} void initializeLayouts() { auto torch_module = THPObjectPtr(PyImport_ImportModule("torch")); if (!torch_module) throw python_error(); - PyObject* strided_layout = - THPLayout_New(at::Layout::Strided, "torch.strided"); - Py_INCREF(strided_layout); - if (PyModule_AddObject(torch_module, "strided", strided_layout) != 0) { - throw python_error(); - } - registerLayoutObject((THPLayout*)strided_layout, at::Layout::Strided); - - PyObject* sparse_coo_layout = - THPLayout_New(at::Layout::Sparse, "torch.sparse_coo"); - Py_INCREF(sparse_coo_layout); - if (PyModule_AddObject(torch_module, "sparse_coo", sparse_coo_layout) != 0) { - throw python_error(); - } - registerLayoutObject((THPLayout*)sparse_coo_layout, at::Layout::Sparse); - - REGISTER_LAYOUT(sparse_csr, SparseCsr) - REGISTER_LAYOUT(sparse_csc, SparseCsc) - REGISTER_LAYOUT(sparse_bsr, SparseBsr) - REGISTER_LAYOUT(sparse_bsc, SparseBsc) - - PyObject* mkldnn_layout = THPLayout_New(at::Layout::Mkldnn, "torch._mkldnn"); - Py_INCREF(mkldnn_layout); - if (PyModule_AddObject(torch_module, "_mkldnn", mkldnn_layout) != 0) { - throw python_error(); - } - registerLayoutObject((THPLayout*)mkldnn_layout, at::Layout::Mkldnn); - - REGISTER_LAYOUT(jagged, Jagged); + registerLayout(torch_module, at::Layout::Strided, "strided", "torch.strided"); + registerLayout( + torch_module, at::Layout::Sparse, "sparse_coo", "torch.sparse_coo"); + registerLayout( + torch_module, at::Layout::SparseCsr, "sparse_csr", "torch.sparse_csr"); + registerLayout( + torch_module, at::Layout::SparseCsc, "sparse_csc", "torch.sparse_csc"); + registerLayout( + torch_module, at::Layout::SparseBsr, "sparse_bsr", "torch.sparse_bsr"); + registerLayout( + torch_module, at::Layout::SparseBsc, "sparse_bsc", "torch.sparse_bsc"); + registerLayout(torch_module, at::Layout::Mkldnn, "_mkldnn", "torch._mkldnn"); + registerLayout(torch_module, at::Layout::Jagged, "jagged", "torch.jagged"); } } // namespace torch::utils diff --git a/torch/csrc/utils/tensor_memoryformats.cpp b/torch/csrc/utils/tensor_memoryformats.cpp index c1a3ff326493a..0ba5c8bd842c5 100644 --- a/torch/csrc/utils/tensor_memoryformats.cpp +++ b/torch/csrc/utils/tensor_memoryformats.cpp @@ -31,14 +31,12 @@ void initializeMemoryFormats() { auto add_memory_format = [&](at::MemoryFormat format, const char* name) { std::string module_name = "torch."; - PyObject* memory_format = THPMemoryFormat_New(format, module_name + name); - Py_INCREF(memory_format); - if (PyModule_AddObject(torch_module, name, memory_format) != 0) { - Py_DECREF(memory_format); + THPObjectPtr memory_format(THPMemoryFormat_New(format, module_name + name)); + if (PyModule_AddObjectRef(torch_module, name, memory_format.get()) != 0) { throw python_error(); } - Py_INCREF(memory_format); - memory_format_registry[static_cast(format)] = memory_format; + memory_format_registry[static_cast(format)] = + memory_format.release(); }; add_memory_format(at::MemoryFormat::Preserve, "preserve_format"); diff --git a/torch/csrc/utils/tensor_new.cpp b/torch/csrc/utils/tensor_new.cpp index 8b758df7224a9..97ae280dfe332 100644 --- a/torch/csrc/utils/tensor_new.cpp +++ b/torch/csrc/utils/tensor_new.cpp @@ -1382,26 +1382,16 @@ static void _validate_sparse_compressed_tensor_args_template( ARG_SIZE, ARGS_COUNT }; - static std::string sig; - switch (required_layout) { - case c10::Layout::SparseCsr: - sig = - "_validate_sparse_csr_tensor(PyObject* crow_indices, PyObject* col_indices, PyObject* values, IntArrayRef size)"; - break; - case c10::Layout::SparseCsc: - sig = - "_validate_sparse_csc_tensor(PyObject* ccol_indices, PyObject* row_indices, PyObject* values, IntArrayRef size)"; - break; - case c10::Layout::SparseBsr: - sig = - "_validate_sparse_bsr_tensor(PyObject* crow_indices, PyObject* col_indices, PyObject* values, IntArrayRef size)"; - break; - case c10::Layout::SparseBsc: - sig = - "_validate_sparse_bsc_tensor(PyObject* ccol_indices, PyObject* row_indices, PyObject* values, IntArrayRef size)"; - break; - default:; - } + constexpr const char* sig = [] { + if constexpr (required_layout == c10::Layout::SparseCsr) + return "_validate_sparse_csr_tensor(PyObject* crow_indices, PyObject* col_indices, PyObject* values, IntArrayRef size)"; + else if constexpr (required_layout == c10::Layout::SparseCsc) + return "_validate_sparse_csc_tensor(PyObject* ccol_indices, PyObject* row_indices, PyObject* values, IntArrayRef size)"; + else if constexpr (required_layout == c10::Layout::SparseBsr) + return "_validate_sparse_bsr_tensor(PyObject* crow_indices, PyObject* col_indices, PyObject* values, IntArrayRef size)"; + else if constexpr (required_layout == c10::Layout::SparseBsc) + return "_validate_sparse_bsc_tensor(PyObject* ccol_indices, PyObject* row_indices, PyObject* values, IntArrayRef size)"; + }(); static PythonArgParser parser({sig}); ParsedArgs parsed_args; @@ -1744,7 +1734,7 @@ Tensor asarray( std::optional dtype, std::optional device, std::optional copy, - bool requires_grad) { + std::optional requires_grad) { Tensor tensor; bool force_copy = copy.value_or(false); @@ -1761,9 +1751,17 @@ Tensor asarray( if (THPVariable_Check(obj)) { tensor = THPVariable_Unpack(obj); } + bool return_requires_grad = + requires_grad.value_or(tensor.defined() ? tensor.requires_grad() : false); + if (return_requires_grad && !requires_grad) { + TORCH_WARN_ONCE( + "torch.asarray: unspecified requires_grad now defaults to obj.requires_grad " + "instead of False. Pass requires_grad=False explicitly to get the old behavior " + "and silence this warning.") + } #ifdef USE_NUMPY - if (is_numpy_available()) { + if (!tensor.defined() && is_numpy_available()) { // Check whether 'obj' is a NumPy Array or Scalar. bool is_numpy_array = PyArray_Check(obj); bool is_numpy_scalar = PyArray_CheckScalar(obj); @@ -1809,7 +1807,8 @@ Tensor asarray( // Check whether 'obj' implements the buffer protocol if (!tensor.defined() && PyObject_CheckBuffer(obj) != 0) { - tensor = tensor_frombuffer(obj, dtype_unwrapped, -1, 0, requires_grad); + tensor = + tensor_frombuffer(obj, dtype_unwrapped, -1, 0, return_requires_grad); } if (tensor.defined()) { @@ -1857,10 +1856,10 @@ Tensor asarray( // Setting 'requires_grad' when the tensor is not a leaf does not work. // Whenever that happens, we have to use 'detach'. - if (!tensor.is_leaf() && !requires_grad) { + if (!tensor.is_leaf() && !return_requires_grad) { tensor = tensor.detach(); } else { - tensor.set_requires_grad(requires_grad); + tensor.set_requires_grad(return_requires_grad); } } else { // Undefined tensor means it does not implement neither DLPack nor @@ -1881,7 +1880,7 @@ Tensor asarray( /* copy_variables = */ false, /* copy_numpy = */ false, /* type_inference = */ !dtype.has_value()); - tensor.set_requires_grad(requires_grad); + tensor.set_requires_grad(return_requires_grad); } return tensor; diff --git a/torch/csrc/utils/tensor_new.h b/torch/csrc/utils/tensor_new.h index 8ae71fcde4cfb..903645c40dc1d 100644 --- a/torch/csrc/utils/tensor_new.h +++ b/torch/csrc/utils/tensor_new.h @@ -132,5 +132,5 @@ at::Tensor asarray( std::optional dtype, std::optional device, std::optional copy, - bool requires_grad); + std::optional requires_grad); } // namespace torch::utils diff --git a/torch/csrc/utils/tensor_qschemes.cpp b/torch/csrc/utils/tensor_qschemes.cpp index f85d091bd57a0..6c86be1ff7f8a 100644 --- a/torch/csrc/utils/tensor_qschemes.cpp +++ b/torch/csrc/utils/tensor_qschemes.cpp @@ -20,13 +20,12 @@ void initializeQSchemes() { for (const auto i : c10::irange(at::COMPILE_TIME_NUM_QSCHEMES)) { auto qscheme = static_cast(i); - PyObject* qscheme_obj = THPQScheme_New(qscheme, toString(qscheme)); - thp_qscheme_array[static_cast(qscheme)] = qscheme_obj; - Py_INCREF(qscheme_obj); - if (PyModule_AddObject( - torch_module, toString(qscheme).c_str(), qscheme_obj) != 0) { + THPObjectPtr qscheme_obj(THPQScheme_New(qscheme, toString(qscheme))); + if (PyModule_AddObjectRef( + torch_module, toString(qscheme).c_str(), qscheme_obj.get()) != 0) { throw python_error(); } + thp_qscheme_array[static_cast(qscheme)] = qscheme_obj.release(); } } diff --git a/torch/cuda/__init__.py b/torch/cuda/__init__.py index 08ac01b4093ec..9b7a6be38941f 100644 --- a/torch/cuda/__init__.py +++ b/torch/cuda/__init__.py @@ -52,7 +52,9 @@ _is_in_bad_fork = getattr(torch._C, "_cuda_isInBadFork", lambda: False) _HAS_PYNVML = False +_HAS_AMDSMI = False _PYNVML_ERR = None +_AMDSMI_ERR = None try: from torch import version as _version @@ -104,8 +106,13 @@ def __enter__(self) -> None: def __exit__(self, type: Any, value: Any, traceback: Any) -> None: ctypes.CDLL = self.original_CDLL # type: ignore[misc] - with _amdsmi_cdll_hook(): - import amdsmi # type: ignore[import] + try: + with _amdsmi_cdll_hook(): + import amdsmi # type: ignore[import] + _HAS_AMDSMI = True + except ModuleNotFoundError as err: + _AMDSMI_ERR = err + raise _HAS_PYNVML = True except ModuleNotFoundError: @@ -883,7 +890,7 @@ def parse_list_with_prefix(lst: str, prefix: str) -> list[str]: def _raw_device_count_amdsmi() -> int: - if not _HAS_PYNVML: # If amdsmi is not available + if not _HAS_AMDSMI: return -1 try: amdsmi.amdsmi_init() @@ -917,7 +924,7 @@ def _raw_device_count_nvml() -> int: def _raw_device_uuid_amdsmi() -> list[str] | None: from ctypes import byref, c_int, c_void_p, CDLL, create_string_buffer - if not _HAS_PYNVML: # If amdsmi is not available + if not _HAS_AMDSMI: return None try: amdsmi.amdsmi_init() @@ -1306,11 +1313,10 @@ def _get_pynvml_handler(device: Device = None): def _get_amdsmi_handler(device: Device = None): - if not _HAS_PYNVML: + if not _HAS_AMDSMI: raise ModuleNotFoundError( "amdsmi does not seem to be installed or it can't be imported." - # pyrefly: ignore [invalid-inheritance] - ) from _PYNVML_ERR + ) from _AMDSMI_ERR try: amdsmi.amdsmi_init() except amdsmi.AmdSmiException as e: @@ -1322,11 +1328,49 @@ def _get_amdsmi_handler(device: Device = None): return handle +_cached_hip_to_amdsmi: dict[int, int] | None = None + + +def _get_amdsmi_device_index_from_hip_index(device: int) -> int: + r"""Return amdsmi index from HIP device index. They are not always the same. + + Assume amdsmi_init() already completes successfully.""" + global _cached_hip_to_amdsmi + if _cached_hip_to_amdsmi is None: + amdsmi_handles = amdsmi.amdsmi_get_processor_handles() + + def gen(): + for amdsmi_idx, handle in enumerate(amdsmi_handles): + info = amdsmi.amdsmi_get_gpu_enumeration_info(handle) + if "hip_id" in info: + yield info["hip_id"], amdsmi_idx + + _cached_hip_to_amdsmi = dict(gen()) + if not _cached_hip_to_amdsmi and len(amdsmi_handles) > 1: + warnings.warn( + "Cannot translate HIP ID to AMD SMI ID due to" + " lack of translation information prior to ROCm 6.4." + " Functions that rely on amdsmi" + " (e.g. temperature()) may operate on wrong devices." + ) + if device not in _cached_hip_to_amdsmi: + warnings.warn( + f"Cannot translate HIP ID {device} to AMD SMI ID due to" + " undetected HIP ID from amdsmi." + " amdsmi_get_gpu_enumeration_info() only report these HIP IDs" + f" {list(_cached_hip_to_amdsmi.keys())}." + " Functions that rely on amdsmi" + " (e.g. temperature()) may operate on wrong devices." + ) + return _cached_hip_to_amdsmi.get(device, device) + + def _get_amdsmi_device_index(device: Device) -> int: r"""Return the amdsmi index of the device, taking visible_devices into account.""" idx = _get_device_index(device, optional=True) visible_devices = _parse_visible_devices() - if type(visible_devices[0]) is str: + visible_device_is_str = type(visible_devices[0]) is str + if visible_device_is_str: uuids = _raw_device_uuid_amdsmi() if uuids is None: raise RuntimeError("Can't get device UUIDs") @@ -1339,7 +1383,10 @@ def _get_amdsmi_device_index(device: Device) -> int: raise RuntimeError( f"device {idx} is not visible (HIP_VISIBLE_DEVICES={visible_devices})" ) - return idx_map[idx] + if visible_device_is_str: + return idx_map[idx] + else: + return _get_amdsmi_device_index_from_hip_index(idx_map[idx]) def _get_amdsmi_device_memory_used(device: Device = None) -> int: @@ -1937,6 +1984,7 @@ def _compile_kernel( "amp", "caching_allocator_alloc", "caching_allocator_delete", + "caching_allocator_disabled", "caching_allocator_enable", "can_device_access_peer", "check_error", diff --git a/torch/cuda/_annotate_cuda_graph_trace.py b/torch/cuda/_annotate_cuda_graph_trace.py new file mode 100644 index 0000000000000..92a0c665a027c --- /dev/null +++ b/torch/cuda/_annotate_cuda_graph_trace.py @@ -0,0 +1,362 @@ +#!/usr/bin/env python3 +"""Post-process a profiler trace to add CUDA graph kernel annotations. + +Reads a profiler trace (gzipped or plain JSON) and a kernel annotations +pickle, matches kernel events by their graph node id, and writes an +annotated trace with the annotation fields added to each kernel event's +args (displayed alongside grid/block size in trace viewers). + +The annotations pickle is auto-discovered from the trace file's parent +directory (one level up, matching the rank from the trace filename). + +Usage: + python -m torch.cuda._annotate_cuda_graph_trace [-a ] [-o ] + +Examples: + # Auto-discover annotations pickle from trace location + python -m torch.cuda._annotate_cuda_graph_trace \\ + traces/step_000000000014/000000.*.pt.trace.json.gz + + # Explicit annotations pickle + python -m torch.cuda._annotate_cuda_graph_trace trace.json.gz -a annotations.pkl +""" + +import argparse +import gzip +import json +import pickle +import re +import sys +from collections import defaultdict +from pathlib import Path +from typing import Any + + +_WORK_CATEGORIES = {"kernel", "gpu_memcpy", "gpu_memset"} + + +def _move_overlapping_to_stream( + trace: dict, default_stream: int = 7, overlap_stream: int = 8 +) -> int: + """Move graphed kernels that overlap with their predecessor to a separate stream. + + Perfetto cannot display overlapping (non-nested) events on the same + stream -- they get hidden. This pass detects graphed kernel events on + *default_stream* whose start timestamp falls before the previous + kernel's end, and moves them to *overlap_stream* so they're visible. + + Returns the number of events moved. + """ + graphed_on_default = [ + e + for e in trace["traceEvents"] + if e.get("cat") == "kernel" + and e.get("tid") == default_stream + and e.get("args", {}).get("graph node id", 0) != 0 + ] + graphed_on_default.sort(key=lambda e: e["ts"]) + + moved = 0 + prev_end = 0.0 + for event in graphed_on_default: + ts = event["ts"] + dur = event.get("dur", 0) + if ts < prev_end: + event["tid"] = overlap_stream + event.get("args", {})["stream"] = overlap_stream + moved += 1 + else: + prev_end = ts + dur + + return moved + + +def _fix_overlapping_timestamps(trace: dict, max_adjust_us: float = 1.0) -> int: + """Clamp graphed kernel/memcpy timestamps so they don't overlap on the same stream. + + CUPTI can produce slightly overlapping timestamps for consecutive graphed + events, causing Perfetto to hide events that sit entirely "under" their + neighbours. This pass sorts graphed work events per stream and ensures + each event starts at or after the previous event's end. + + Overlaps larger than *max_adjust_us* are flagged as warnings and left + unchanged, since they likely indicate a real issue rather than CUPTI + timestamp jitter. + + Returns the number of events adjusted. + """ + per_stream: dict[int, list[dict]] = defaultdict(list) + for event in trace["traceEvents"]: + if ( + event.get("cat") in _WORK_CATEGORIES + and event.get("args", {}).get("graph node id", 0) != 0 + ): + per_stream[event.get("tid")].append(event) + + adjusted = 0 + for tid, events in per_stream.items(): + events.sort(key=lambda e: e["ts"]) + prev_end = 0.0 + for event in events: + ts = event["ts"] + dur = event.get("dur", 0) + if ts < prev_end: + overlap = prev_end - ts + if overlap > max_adjust_us: + print( + f"WARNING: large overlap {overlap:.3f}us on stream {tid} " + f"for {event.get('name', '?')[:60]}, skipping adjustment" + ) + else: + event["ts"] = prev_end + adjusted += 1 + prev_end = event["ts"] + dur + + return adjusted + + +def annotate_trace( + trace: dict, + annotations: dict[int, list[Any]], + default_stream: int = 7, +) -> int: + """Add annotation fields to kernel events matching the annotations dict. + + Each annotation entry is a list (from nested ``mark_kernels`` scopes). + Fields from all annotations are merged into the event args; if multiple + annotations define the same key, later entries in the list win. + + For graphed events (graph_node_id != 0), reassigns ``tid`` and + ``args["stream"]`` to the stream recorded in annotations, or to + *default_stream* if there is no annotation. Also moves the + corresponding ``ac2g`` flow-finish events to the new tid so that + CPU-to-GPU correlation arrows are preserved. + + Removes ``gpu_user_annotation`` events and orphaned ``ac2g`` events + from streams that have no kernel or memcpy events after reassignment, + since CUPTI replicates these onto every stream during graph replay. + + Returns the number of events annotated. + """ + # Build an index of ac2g 'f' events keyed by (tid, ts) so we can + # move them together with the kernel events they correspond to. + ac2g_f_index: dict[tuple, list] = {} + for event in trace["traceEvents"]: + if event.get("cat") == "ac2g" and event.get("ph") == "f": + key = (event.get("tid"), event.get("ts")) + ac2g_f_index.setdefault(key, []).append(event) + + annotated = 0 + for event in trace.get("traceEvents", []): + args = event.get("args", {}) + graph_node_id = args.get("graph node id") + if graph_node_id is None or graph_node_id == 0: + continue + stream_id = None + if graph_node_id in annotations: + for ann in annotations[graph_node_id]: + if isinstance(ann, dict): + for key, value in ann.items(): + args[key] = str(value) + if "stream" in ann: + stream_id = int(ann["stream"]) + else: + args["annotation"] = str(ann) + annotated += 1 + + # Reassign stream: use annotated stream if available, else default + if stream_id is None: + stream_id = default_stream + old_key = (event.get("tid"), event.get("ts")) + event["tid"] = stream_id + args["stream"] = stream_id + + # Move the matching ac2g 'f' event(s) to the same new tid + for ac2g_event in ac2g_f_index.get(old_key, ()): + ac2g_event["tid"] = stream_id + + # Remove gpu_user_annotation events and ac2g flow-finish events from + # streams that have no real kernel/memcpy/memset work -- these are + # noise replicated by CUPTI onto every stream during graph replay. + tids_with_work = set() + for event in trace["traceEvents"]: + if event.get("cat") in _WORK_CATEGORIES: + tids_with_work.add(event.get("tid")) + + def _is_noise(event: dict) -> bool: + cat = event.get("cat") + if cat == "gpu_user_annotation": + return event.get("tid") not in tids_with_work + if cat == "ac2g" and event.get("ph") == "f": + return event.get("tid") not in tids_with_work + return False + + original_count = len(trace["traceEvents"]) + trace["traceEvents"] = [ + event for event in trace["traceEvents"] if not _is_noise(event) + ] + removed = original_count - len(trace["traceEvents"]) + if removed: + print(f"Removed {removed} noise events from empty streams") + + # Clean up metadata: remove thread_name / thread_sort_index entries + # for noise streams that have no real (non-metadata) events, and add + # thread_name entries for our new annotation streams. + all_tids_in_trace = { + e.get("tid") for e in trace["traceEvents"] if e.get("ph") != "M" + } + # Find the GPU process pid from existing thread_name metadata + gpu_pid = 0 + for event in trace["traceEvents"]: + if ( + event.get("ph") == "M" + and event.get("name") == "thread_name" + and str(event.get("args", {}).get("name", "")).startswith("stream ") + ): + gpu_pid = event.get("pid", 0) + break + + # Remove metadata entries for tids with no non-metadata events + trace["traceEvents"] = [ + event + for event in trace["traceEvents"] + if event.get("ph") != "M" or event.get("tid") in all_tids_in_trace + ] + + # Add thread_name metadata for new annotation tids that lack one + existing_thread_names = { + e.get("tid") + for e in trace["traceEvents"] + if e.get("ph") == "M" and e.get("name") == "thread_name" + } + for tid in sorted(tids_with_work - existing_thread_names): + trace["traceEvents"].append( + { + "ph": "M", + "pid": gpu_pid, + "tid": tid, + "name": "thread_name", + "args": {"name": f"stream {tid}"}, + } + ) + + return annotated + + +def load_trace(path: Path) -> dict: + if path.suffix == ".gz" or path.name.endswith(".json.gz"): + with gzip.open(path, "rt") as f: + return json.load(f) + else: + with open(path) as f: + return json.load(f) + + +def save_trace(trace: dict, path: Path) -> None: + if path.suffix == ".gz" or path.name.endswith(".json.gz"): + with gzip.open(path, "wt") as f: + json.dump(trace, f) + else: + with open(path, "w") as f: + json.dump(trace, f) + + +def _find_annotations_pkl(trace_file: Path) -> Path | None: + """Auto-discover the annotations pickle from the trace file location. + + Trace files live in e.g. ``traces/step_000000000014/000000..pt.trace.json.gz`` + where the leading digits are the rank. The pickle lives one level up: + ``traces/kernel_annotations_rank0_*.pkl``. + """ + match = re.match(r"^(\d+)", trace_file.name) + if not match: + return None + rank = int(match.group(1)) + + traces_dir = trace_file.parent.parent + candidates = sorted(traces_dir.glob(f"kernel_annotations_rank{rank}_*.pkl")) + if candidates: + return candidates[0] + return None + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Annotate a profiler trace with CUDA graph kernel annotations." + ) + parser.add_argument( + "trace_file", type=Path, help="Input trace file (.json or .json.gz)" + ) + parser.add_argument( + "-a", + "--annotations", + type=Path, + default=None, + help="Kernel annotations pickle file. Auto-discovered if omitted.", + ) + parser.add_argument( + "-o", + "--output", + type=Path, + default=None, + help="Output file path. Defaults to .annotated.", + ) + parser.add_argument( + "--default-stream", + type=int, + default=7, + help="Stream ID to assign to unannotated graphed events (default: 7).", + ) + args = parser.parse_args() + + annotations_pkl = args.annotations + if annotations_pkl is None: + annotations_pkl = _find_annotations_pkl(args.trace_file) + if annotations_pkl is None: + print( + f"Could not auto-discover annotations pickle for {args.trace_file}. " + f"Use -a to specify it explicitly.", + file=sys.stderr, + ) + sys.exit(1) + print(f"Auto-discovered annotations: {annotations_pkl}") + + with open(annotations_pkl, "rb") as f: + annotations = pickle.load(f) + print(f"Loaded {len(annotations)} kernel annotations") + + trace = load_trace(args.trace_file) + total_events = len(trace.get("traceEvents", [])) + print(f"Loaded trace with {total_events} events") + + count = annotate_trace(trace, annotations, default_stream=args.default_stream) + print(f"Annotated {count} kernel events") + + overlap_moved = _move_overlapping_to_stream( + trace, default_stream=args.default_stream + ) + if overlap_moved: + print(f"Moved {overlap_moved} overlapping events to stream 8") + + ts_fixed = _fix_overlapping_timestamps(trace) + if ts_fixed: + print(f"Fixed {ts_fixed} overlapping graphed event timestamps") + + output = args.output + if output is None: + name = args.trace_file.name + if name.endswith(".json.gz"): + output = args.trace_file.with_name( + name.replace(".json.gz", ".annotated.json.gz") + ) + elif name.endswith(".json"): + output = args.trace_file.with_suffix(".annotated.json") + else: + output = args.trace_file.with_suffix(args.trace_file.suffix + ".annotated") + + save_trace(trace, output) + print(f"Saved annotated trace to {output}") + + +if __name__ == "__main__": + main() diff --git a/torch/cuda/_graph_annotations.py b/torch/cuda/_graph_annotations.py new file mode 100644 index 0000000000000..16a895686cf41 --- /dev/null +++ b/torch/cuda/_graph_annotations.py @@ -0,0 +1,434 @@ +"""Annotate CUDA graph kernel nodes during capture. + +During CUDA graph capture, ``mark_kernels`` uses ``cudaGraphGetNodes`` +to count nodes before and after the wrapped region. Nodes at indices +``[before, after)`` are the ones added within the scope. Each kernel +or memcpy node found is annotated by its ``toolsId`` so it can later +be matched to profiler trace events. + +The annotations can be pickled and later merged into a Chrome profiler +trace using ``torch.cuda._annotate_cuda_graph_trace``. + +Requires ``cuda.bindings`` package and a CUDA driver that supports +``cudaGraphNodeGetToolsId`` (CUDA >= 13.1 or appropriate cuda-compat). +When unavailable, ``mark_kernels`` silently becomes a no-op. + +Usage during capture:: + + from torch.cuda._graph_annotations import ( + enable_annotations, + mark_kernels, + resolve_pending_annotations, + remap_to_exec_graph, + ) + + enable_annotations() + + with torch.cuda.graph(graph): + with mark_kernels("phase_A"): + y = workload_a(x) + with mark_kernels("phase_B"): + z = workload_b(y) + resolve_pending_annotations() + + remap_to_exec_graph(graph) +""" + +from collections import defaultdict +from contextlib import contextmanager +from logging import getLogger +from typing import Any + +import torch +from torch.cuda._utils import _check_cuda_bindings, _HAS_CUDA_BINDINGS + + +try: + from cuda.bindings import ( # pyrefly: ignore[missing-import] + runtime as _cuda_runtime, + ) +except ImportError: + _cuda_runtime = None # type: ignore[assignment] + + +logger = getLogger(__name__) + + +# Tri-state: None = not probed, True = available, False = unavailable. +# Deferred to first use to avoid premature CUDA initialization. +_tools_id_available: bool | None = None + +# Global kill switch. When False, mark_kernels and mark_stream are no-ops. +_annotations_enabled: bool = False + + +def enable_annotations() -> None: + """Enable kernel annotation recording.""" + global _annotations_enabled + _annotations_enabled = True + + +def disable_annotations() -> None: + """Disable kernel annotation recording.""" + global _annotations_enabled + _annotations_enabled = False + + +def _is_tools_id_unavailable() -> bool: + """Return True if we already know cudaGraphNodeGetToolsId is missing.""" + if not _HAS_CUDA_BINDINGS: + return True + if _tools_id_available is False: + return True + if not hasattr(_cuda_runtime, "cudaGraphNodeGetToolsId"): + return True + return False + + +def _get_tools_id(node: Any) -> int | None: + """Return the toolsId for a graph node, or None if unavailable.""" + global _tools_id_available + if _tools_id_available is None: + try: + tools_id = _check_cuda_bindings( + _cuda_runtime.cudaGraphNodeGetToolsId( # pyrefly: ignore[missing-attribute] + node + ) + ) + except Exception: + _tools_id_available = False + logger.info( + "cudaGraphNodeGetToolsId not available; " + "CUDA graph kernel annotations will be disabled" + ) + return None + _tools_id_available = True + return tools_id + return _check_cuda_bindings( + _cuda_runtime.cudaGraphNodeGetToolsId( # pyrefly: ignore[missing-attribute] + node + ) + ) + + +def _get_capture_graph(stream: Any) -> Any: + """Return the graph handle for the active capture, or None.""" + status, _id, graph, _deps, _edge_data, _num_deps = _check_cuda_bindings( + _cuda_runtime.cudaStreamGetCaptureInfo( # pyrefly: ignore[missing-attribute] + stream + ) + ) + if ( + status + != _cuda_runtime.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive # pyrefly: ignore[missing-attribute] + ): + return None + return graph + + +def _get_node_count(graph: Any) -> int: + """Return the number of nodes currently in the graph.""" + _, num = _check_cuda_bindings( + _cuda_runtime.cudaGraphGetNodes( # pyrefly: ignore[missing-attribute] + graph, numNodes=0 + ) + ) + return num + + +# toolsId -> list of annotation objects. +_kernel_annotations: defaultdict[int, list[Any]] = defaultdict(list) + +# Node types we annotate. Initialized lazily to avoid touching cuda.bindings +# at import time. +_ANNOTATABLE_TYPES: set[Any] | None = None + + +def _get_annotatable_types() -> set[Any]: + global _ANNOTATABLE_TYPES + if _ANNOTATABLE_TYPES is None: + _ANNOTATABLE_TYPES = { + _cuda_runtime.cudaGraphNodeType.cudaGraphNodeTypeKernel, # pyrefly: ignore[missing-attribute] + _cuda_runtime.cudaGraphNodeType.cudaGraphNodeTypeMemcpy, # pyrefly: ignore[missing-attribute] + } + return _ANNOTATABLE_TYPES + + +# Pending scopes: (annotation, start_node_index, end_node_index). +_pending_scopes: list[tuple[Any, int, int]] = [] + +# Graph handle saved during capture for post-capture resolution. +_capture_graph: Any = None + +# Capture graph ID saved by resolve_pending_annotations for remap_to_exec_graph. +_last_capture_graph_id: int | None = None + + +@contextmanager # type: ignore[arg-type] +def mark_kernels(annotation: str | dict[str, Any]): + """Context manager that records node index ranges for later annotation. + + During capture, calls ``cudaGraphGetNodes`` to count graph nodes before + and after the scope. Nodes at indices ``[before, after)`` were added + inside the scope. After capture, ``resolve_pending_annotations`` + enumerates all nodes and annotates kernel/memcpy nodes in those ranges. + + Must be called inside an active ``torch.cuda.graph()`` capture. If the + current stream is not capturing, or if ``cudaGraphNodeGetToolsId`` is not + available, the context manager is a no-op. + + Args: + annotation: Arbitrary object appended to the annotation list for + every kernel/memcpy node whose index falls within this scope. + """ + if not _annotations_enabled or _is_tools_id_unavailable(): + yield + return + + if isinstance(annotation, str): + annotation = {"str": annotation} + + stream = _cuda_runtime.cudaStream_t( # pyrefly: ignore[missing-attribute] + init_value=torch.cuda.current_stream().cuda_stream + ) + graph = _get_capture_graph(stream) + if graph is None: + yield + return + + global _capture_graph + _capture_graph = graph + + start_count = _get_node_count(graph) + + yield + + end_count = _get_node_count(graph) + + if end_count > start_count: + _pending_scopes.append((annotation, start_count, end_count)) + + +def resolve_pending_annotations() -> None: + """Resolve pending scope index ranges into kernel annotations. + + Enumerates all graph nodes and annotates kernel/memcpy nodes whose + indices fall within recorded scope ranges. Must be called while still + inside the ``torch.cuda.graph()`` capture context. + """ + global _capture_graph + if not _pending_scopes: + _capture_graph = None + return + + # Get a fresh graph handle from the active capture. + stream = _cuda_runtime.cudaStream_t( # pyrefly: ignore[missing-attribute] + init_value=torch.cuda.current_stream().cuda_stream + ) + graph = _get_capture_graph(stream) + if graph is None: + graph = _capture_graph + if graph is None: + logger.warning("resolve_pending_annotations: no graph handle available") + _pending_scopes.clear() + return + + try: + num = _get_node_count(graph) + if num == 0: + _pending_scopes.clear() + _capture_graph = None + return + + nodes, num = _check_cuda_bindings( + _cuda_runtime.cudaGraphGetNodes( # pyrefly: ignore[missing-attribute] + graph, numNodes=num + ) + ) + + # Save capture graph ID for remap_to_exec_graph. + global _last_capture_graph_id + if num > 0: + first_tid = _get_tools_id(nodes[0]) + _last_capture_graph_id = (first_tid >> 32) if first_tid else None + + annotatable = _get_annotatable_types() + + # Sort by (start, -end, -append_index). The append index encodes + # nesting depth: inner context managers exit first, so they are + # appended to _pending_scopes first (smaller index). Using + # -append_index as tiebreaker ensures that for same-range scopes + # the outer scope (larger index) sorts first and is pushed onto + # the stack first, leaving the inner scope on top. + sorted_scopes = sorted( + ( + (ann, start, end, i) + for i, (ann, start, end) in enumerate(_pending_scopes) + ), + key=lambda s: (s[1], -s[2], -s[3]), + ) + scope_ptr = 0 + active_stack: list[tuple[int, Any]] = [] # (end_idx, annotation) + + for i in range(num): + # Pop scopes whose range ended. + while active_stack and active_stack[-1][0] <= i: + active_stack.pop() + + # Push scopes that start at or before this index. + while scope_ptr < len(sorted_scopes) and sorted_scopes[scope_ptr][1] <= i: + ann, _start_idx, end_idx, _idx = sorted_scopes[scope_ptr] + if end_idx > i: + active_stack.append((end_idx, ann)) + scope_ptr += 1 + + if not active_stack: + continue + + node = nodes[i] + node_type = _check_cuda_bindings( + _cuda_runtime.cudaGraphNodeGetType( # pyrefly: ignore[missing-attribute] + node + ) + ) + if node_type not in annotatable: + continue + + tools_id = _get_tools_id(node) + if tools_id is None: + logger.warning( + "resolve_pending_annotations: toolsId unavailable, aborting" + ) + _pending_scopes.clear() + _capture_graph = None + return + + if len(active_stack) == 1: + _kernel_annotations[tools_id].append(active_stack[0][1]) + else: + # Merge all active scopes into one dict. Inner scopes sit + # on top of the stack. Iterating reversed (inner first) + # with setdefault lets the inner scope's values win for + # overlapping keys (e.g. name, stream) while outer scopes + # fill in any missing keys. + merged: dict[str, Any] = {} + for _, ann in reversed(active_stack): + if isinstance(ann, dict): + for ak, av in ann.items(): + merged.setdefault(ak, av) + else: + merged.setdefault("name", ann) + _kernel_annotations[tools_id].append(merged) + except Exception: + logger.exception("resolve_pending_annotations failed") + finally: + _pending_scopes.clear() + _capture_graph = None + + +def remap_to_exec_graph(torch_cuda_graph: torch.cuda.CUDAGraph) -> None: + """Remap annotation keys from capture graph ID to exec graph ID. + + During capture, toolsId encodes the capture graph's ID in the upper + 32 bits. After instantiation, the profiler uses the exec graph's ID. + This function rewrites the keys so annotations match the trace. + + Must be called after the ``torch.cuda.graph()`` context exits. + """ + if not _kernel_annotations: + return + + exec_handle = _cuda_runtime.cudaGraphExec_t( # pyrefly: ignore[missing-attribute] + init_value=torch_cuda_graph.raw_cuda_graph_exec() + ) + exec_graph_id = _check_cuda_bindings( + _cuda_runtime.cudaGraphExecGetId( # pyrefly: ignore[missing-attribute] + exec_handle + ) + ) + + # Only remap annotations from the most recent capture graph. + # Previously remapped annotations (from earlier captures) keep their + # correct exec graph IDs. + capture_graph_id = _last_capture_graph_id + remapped: dict[int, list[Any]] = {} + for tools_id, ann_list in _kernel_annotations.items(): + graph_id = tools_id >> 32 + if capture_graph_id is not None and graph_id != capture_graph_id: + # Belongs to a different graph — keep as-is. + remapped[tools_id] = ann_list + continue + node_id = tools_id & 0xFFFFFFFF + new_tools_id = (exec_graph_id << 32) | node_id + if new_tools_id in remapped: + remapped[new_tools_id].extend(ann_list) + else: + remapped[new_tools_id] = list(ann_list) + + _kernel_annotations.clear() + _kernel_annotations.update(remapped) + + +def get_kernel_annotations() -> dict[int, list[Any]]: + """Return the current kernel annotation map (toolsId -> annotations).""" + return _kernel_annotations + + +def clear_kernel_annotations() -> None: + """Clear all recorded kernel annotations and pending scopes.""" + global _capture_graph + _kernel_annotations.clear() + _pending_scopes.clear() + _capture_graph = None + + +# Counter-based stream ID registry. IDs start at 60 (above the highest +# observed non-graphed CUDA stream ID) so every assigned lane is visually +# distinct in Perfetto and doesn't collide with real streams. +_stream_id_counter: int = 60 +_stream_id_map: dict[int, int] = {} + + +def _get_stream_id(stream: torch.cuda.Stream) -> int: + """Return a small, stable stream ID for the given CUDA stream.""" + global _stream_id_counter + key = stream.cuda_stream + if key not in _stream_id_map: + _stream_id_map[key] = _stream_id_counter + _stream_id_counter += 1 + return _stream_id_map[key] + + +def get_stream_for_pg(pg_key: str) -> int: + """Return a unique stream ID for the given process group key.""" + global _stream_id_counter + if pg_key not in _stream_id_map: + _stream_id_map[pg_key] = _stream_id_counter # type: ignore[assignment] + _stream_id_counter += 1 + return _stream_id_map[pg_key] # type: ignore[return-value] + + +@contextmanager # type: ignore[arg-type] +def mark_stream(stream: torch.cuda.Stream, annotation: str | dict[str, Any]): + """Switch to stream, inject its ID into annotation, and mark kernels. + + If *stream* is already the current stream, no stream switch or stream ID + injection happens — the kernels stay on whatever stream is active (which + keeps the trace faithful when e.g. FSDP uses the current stream for + copy-in instead of a separate one). + """ + if not _annotations_enabled: + with torch.cuda.stream(stream): + yield + return + if stream.cuda_stream == torch.cuda.current_stream().cuda_stream: + with mark_kernels(annotation): + yield + else: + if isinstance(annotation, str): + annotation = {"str": annotation} + if isinstance(annotation, dict): + annotation["stream"] = _get_stream_id(stream) + with torch.cuda.stream(stream): + with mark_kernels(annotation): + yield diff --git a/torch/cuda/_utils.py b/torch/cuda/_utils.py index d842e8b56ef41..8d9501afcbff3 100644 --- a/torch/cuda/_utils.py +++ b/torch/cuda/_utils.py @@ -4,6 +4,17 @@ import torch + +try: + from cuda.bindings import ( # pyrefly: ignore[missing-import] + runtime as _cuda_bindings_runtime, + ) + + _HAS_CUDA_BINDINGS = True +except ImportError: + _cuda_bindings_runtime = None # type: ignore[assignment] + _HAS_CUDA_BINDINGS = False + # The _get_device_index has been moved to torch.utils._get_device_index from torch._utils import _get_device_index as _torch_get_device_index @@ -59,6 +70,35 @@ def _check_cuda(result: int) -> None: raise RuntimeError(f"CUDA error: {error_message}") +def _check_cuda_bindings(result: Any) -> Any: + """Check a cuda.bindings (cuda-python) call result for errors. + + All cuda.bindings runtime calls return ``(error, *outputs)``. This + helper unpacks the tuple, raises on non-success, and returns the + outputs (``None`` for zero outputs, scalar for one, tuple otherwise). + """ + if not _HAS_CUDA_BINDINGS: + raise RuntimeError("cuda.bindings is not available") + err, *out = result + if ( + err + != _cuda_bindings_runtime.cudaError_t.cudaSuccess # pyrefly: ignore[missing-attribute] + ): + _, err_str = ( + _cuda_bindings_runtime.cudaGetErrorString( # pyrefly: ignore[missing-attribute] + err + ) + ) + if isinstance(err_str, bytes): + err_str = err_str.decode() + raise RuntimeError(f"CUDA error: {err} ({err_str})") + if len(out) == 0: + return None + if len(out) == 1: + return out[0] + return out + + def _get_hiprtc_library() -> ctypes.CDLL: try: # pyrefly: ignore [import-error, missing-import] diff --git a/torch/cuda/graphs.py b/torch/cuda/graphs.py index 43d559ddc2440..91bbc99a1dee3 100644 --- a/torch/cuda/graphs.py +++ b/torch/cuda/graphs.py @@ -197,6 +197,12 @@ class graph: may be unsafe. "global" will error on actions in other threads, "thread_local" will only error for actions in the current thread, and "relaxed" will not error on actions. Do NOT change this setting unless you're familiar with `cudaStreamCaptureMode `_ + enable_annotations (bool, optional): If ``True``, enables kernel annotation + recording on entry and automatically calls + :func:`~torch.cuda._graph_annotations.resolve_pending_annotations` before + the capture ends. Annotations are **not** cleared on exit so that multiple + graphs in the same workload can accumulate annotations. + Requires ``cuda.bindings`` package and cuda-compat >= 13.1 or CUDA driver >= 13.1. .. note:: For effective memory sharing, if you pass a ``pool`` used by a previous capture and the previous capture @@ -217,6 +223,7 @@ def __init__( pool: _POOL_HANDLE | None = None, stream: torch.cuda.Stream | None = None, capture_error_mode: str = "global", + enable_annotations: bool = False, ): # Lazy-init of default_capture_stream helps avoid circular-import errors. # Not thread safe, but graphs already have the general (explicitly documented) @@ -233,6 +240,7 @@ def __init__( self.stream_ctx = torch.cuda.stream(self.capture_stream) self.cuda_graph = cuda_graph self.capture_error_mode = capture_error_mode + self._enable_annotations = enable_annotations def __enter__(self) -> None: # Free as much memory as we can for the graph @@ -250,6 +258,11 @@ def __enter__(self) -> None: # pyrefly: ignore [missing-attribute] torch._C._host_emptyCache() + if self._enable_annotations: + from torch.cuda._graph_annotations import enable_annotations as _enable_ann + + _enable_ann() + # Stackoverflow seems comfortable with this pattern # https://stackoverflow.com/questions/26635684/calling-enter-and-exit-manually#39172487 self.stream_ctx.__enter__() @@ -262,8 +275,18 @@ def __enter__(self) -> None: ) def __exit__(self, *args: object) -> None: + if self._enable_annotations: + from torch.cuda._graph_annotations import resolve_pending_annotations + + resolve_pending_annotations() + self.cuda_graph.capture_end() self.stream_ctx.__exit__(*args) + + if self._enable_annotations: + from torch.cuda._graph_annotations import remap_to_exec_graph + + remap_to_exec_graph(self.cuda_graph) # returning None should propagate exceptions from either capture_end or stream_ctx.__exit__() diff --git a/torch/cuda/green_contexts.py b/torch/cuda/green_contexts.py index c08d2c061e404..c7aab1b749db1 100644 --- a/torch/cuda/green_contexts.py +++ b/torch/cuda/green_contexts.py @@ -1,6 +1,10 @@ import torch +__all__ = [ + "GreenContext", +] + _GreenContext = object SUPPORTED = False @@ -19,18 +23,59 @@ class GreenContext(_GreenContext): """ @staticmethod - def create(num_sms: int, device_id: int = 0) -> _GreenContext: + def create( + *, + num_sms: int | None = None, + workqueue_scope: str | None = None, + workqueue_concurrency_limit: int | None = None, + device_id: int | None = None, + ) -> _GreenContext: r"""Create a CUDA green context. + At least one of ``num_sms`` or ``workqueue_scope`` must be specified. + Both can be combined to partition SMs and configure workqueues in the + same green context. + Arguments: - num_sms (int): The number of SMs to use in the green context. + num_sms (int, optional): The number of SMs to use in the green + context. When ``None``, SMs are not partitioned. + workqueue_scope (str, optional): Workqueue sharing scope. One of + ``"device_ctx"`` (shared across all contexts, default driver + behaviour) or ``"balanced"`` (non-overlapping workqueues with + other balanced green contexts). When ``None``, no workqueue + configuration is applied. + workqueue_concurrency_limit (int, optional): Maximum number of + concurrent stream-ordered workloads for the workqueue. Requires + ``workqueue_scope`` to be set. device_id (int, optional): The device index of green context. + When ``None``, the current device is used. + """ + if not SUPPORTED: + raise RuntimeError("PyTorch was not built with Green Context support!") + return _GreenContext.create( # type: ignore[attr-defined] + device_id=device_id, + num_sms=num_sms, + workqueue_scope=workqueue_scope, + workqueue_concurrency_limit=workqueue_concurrency_limit, + ) + + @staticmethod + def max_workqueue_concurrency(device_id: int | None = None) -> int: + r"""Return the maximum workqueue concurrency limit for the device. + + This queries the device for the default number of concurrent + stream-ordered workloads supported by workqueue configuration + resources. + + Arguments: + device_id (int, optional): The device index to query. When + ``None``, the current device is used. """ if not SUPPORTED: raise RuntimeError("PyTorch was not built with Green Context support!") - return _GreenContext.create(num_sms, device_id) # type: ignore[attr-defined] + return _GreenContext.max_workqueue_concurrency(device_id=device_id) # type: ignore[attr-defined] - # Note that these functions are bypassed by we define them here + # Note that these functions are bypassed but we define them here # for Sphinx documentation purposes def set_context(self) -> None: # pylint: disable=useless-parent-delegation r"""Make the green context the current context.""" @@ -42,6 +87,6 @@ def pop_context(self) -> None: # pylint: disable=useless-parent-delegation """ return super().pop_context() # type: ignore[misc] - def Stream(self) -> torch.Stream: + def Stream(self) -> "torch.cuda.Stream": r"""Return the CUDA Stream used by the green context.""" return super().Stream() diff --git a/torch/cuda/memory.py b/torch/cuda/memory.py index f8d4daf1dd441..90ab843c139f4 100644 --- a/torch/cuda/memory.py +++ b/torch/cuda/memory.py @@ -32,6 +32,7 @@ __all__ = [ "caching_allocator_alloc", "caching_allocator_delete", + "caching_allocator_disabled", "caching_allocator_enable", "get_per_process_memory_fraction", "set_per_process_memory_fraction", @@ -98,15 +99,6 @@ def _host_allocator(): return torch._C._cuda_cudaHostAllocator() -@contextlib.contextmanager -def _free_mutex(): - torch._C._cuda_lock_mutex() - try: - yield - finally: - torch._C._cuda_unlock_mutex() - - def caching_allocator_alloc(size, device: "Device" = None, stream=None): r"""Perform a memory allocation using the CUDA memory allocator. @@ -166,6 +158,17 @@ def caching_allocator_enable(value: bool = True) -> None: torch._C._cuda_cudaCachingAllocator_enable(value) +@contextlib.contextmanager +def caching_allocator_disabled(): + r"""Context manager that temporarily disables the CUDA caching allocator.""" + prev = torch._C._cuda_cudaCachingAllocator_is_enabled() + caching_allocator_enable(False) + try: + yield + finally: + caching_allocator_enable(prev) + + def set_per_process_memory_fraction(fraction, device: "Device" = None) -> None: r"""Set memory fraction for a process. @@ -277,6 +280,8 @@ def memory_stats(device: "Device" = None) -> dict[str, Any]: cuMemMap and cudaMalloc. - ``"num_device_free"``: number of CUDA free calls. This includes both cuMemUnmap and cudaFree. + - ``"num_oom_rejections"``: number of allocations preemptively rejected by the + throw_on_cudamalloc_oom + per_process_memory_fraction policy. The caching allocator can be configured via ENV to not split blocks larger than a defined size (see Memory Management section of the Cuda Semantics documentation). diff --git a/torch/distributed/__init__.py b/torch/distributed/__init__.py index 24add6a4d8aeb..4f24b45494e6e 100644 --- a/torch/distributed/__init__.py +++ b/torch/distributed/__init__.py @@ -169,7 +169,8 @@ def breakpoint(rank: int = 0, skip: int = 0, timeout_s=3600): # stubs as necessary. # We cannot define stubs directly because they confuse pyre - class _ProcessGroupStub: + class _Stub: pass - sys.modules["torch.distributed"].ProcessGroup = _ProcessGroupStub # type: ignore[attr-defined] + sys.modules["torch.distributed"].GroupName = _Stub # type: ignore[attr-defined] + sys.modules["torch.distributed"].ProcessGroup = _Stub # type: ignore[attr-defined] diff --git a/torch/distributed/_functional_collectives.py b/torch/distributed/_functional_collectives.py index c2a74a5aa1ee7..733147622ad85 100644 --- a/torch/distributed/_functional_collectives.py +++ b/torch/distributed/_functional_collectives.py @@ -204,8 +204,6 @@ def all_gather_tensor( :: N.B. If you pass a PG or a 1D list to perform a MPMD collective, the compiler won't be able to recover that information and perform collective algebraic optimization. Use other forms of input for that. """ - if not self.is_contiguous(): - raise AssertionError("Tensor must be contiguous for all_gather_tensor") group = _resolve_group(group, tag) group_size = c10d._get_group_size_by_name(group) tensor = torch.ops._c10d_functional.all_gather_into_tensor( @@ -1214,6 +1212,8 @@ def _resolve_group( raise AssertionError( "Only 1D mesh is supported, pass in (DeviceMesh, int) together if mesh > 1D" ) + if dist.config.compile_on_one_rank: + return torch.ops._dtensor.mesh_get_process_group(group, 0) return group._dim_group_names[0] elif isinstance(group, tuple): if ( @@ -1223,6 +1223,8 @@ def _resolve_group( ): dmesh = group[0] dim = group[1] + if dist.config.compile_on_one_rank: + return torch.ops._dtensor.mesh_get_process_group(dmesh, dim) return dmesh._dim_group_names[dim] else: raise ValueError( @@ -1832,6 +1834,8 @@ def _group_or_group_name( ) -> dist.ProcessGroup | c10d.GroupName: if isinstance(group, str): return group + elif dist.config.compile_on_one_rank: + return group else: return group.group_name diff --git a/torch/distributed/_ops/device_mesh.py b/torch/distributed/_ops/device_mesh.py index fb431cad79efc..61d4312c3bb04 100644 --- a/torch/distributed/_ops/device_mesh.py +++ b/torch/distributed/_ops/device_mesh.py @@ -62,3 +62,44 @@ def _runtime_compute_coordinate_on_dim_impl(full_mesh: torch.Tensor, index: int) if mesh_coords is None: raise AssertionError return mesh_coords[index] + + +def _get_flattened_submesh_impl(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh: + from torch.distributed.tensor._redistribute import ( + _get_flattened_mesh_by_layout_impl, + ) + + result = _get_flattened_mesh_by_layout_impl(mesh, tuple(mesh_dims)) + if result is None: + raise ValueError(f"No flattened mesh found for mesh_dims={mesh_dims} on {mesh}") + return result + + +@torch.library.custom_op("device_mesh::_get_flattened_submesh", mutates_args=()) +def _get_flattened_submesh(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh: + return _get_flattened_submesh_impl(mesh, mesh_dims) + + +@_get_flattened_submesh.register_fake +def _get_flattened_submesh_fake(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh: + return _get_flattened_submesh_impl(mesh, mesh_dims) + + +def _get_submesh_impl(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh: + all_dim_names = mesh._mesh_dim_names + if all_dim_names is None: + raise ValueError(f"Cannot slice mesh without dim names: {mesh}") + dim_names = tuple(all_dim_names[i] for i in mesh_dims) + if len(dim_names) == 1: + return mesh[dim_names[0]] + return mesh[dim_names] + + +@torch.library.custom_op("device_mesh::_get_submesh", mutates_args=()) +def _get_submesh(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh: + return _get_submesh_impl(mesh, mesh_dims) + + +@_get_submesh.register_fake +def _get_submesh_fake(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh: + return _get_submesh_impl(mesh, mesh_dims) diff --git a/torch/distributed/_symmetric_memory/__init__.py b/torch/distributed/_symmetric_memory/__init__.py index 7052519061130..1a6c19b85385e 100644 --- a/torch/distributed/_symmetric_memory/__init__.py +++ b/torch/distributed/_symmetric_memory/__init__.py @@ -2008,7 +2008,9 @@ def is_nvshmem_available() -> bool: r""" is_nvshmem_available() -> bool - Check if NVSHMEM is available in current build and on current system. + Check if NVSHMEM (CUDA) or rocSHMEM (ROCm) is available in the current + build and usable at runtime. On ROCm, rocSHMEM ``VERSION`` must be at + least 3.3.0 (see ``rocshmem/rocshmem.hpp``). """ try: from torch._C._distributed_c10d import _is_nvshmem_available @@ -2195,8 +2197,100 @@ def wait_signal(hdl: _SymmetricMemory, peer: int) -> None: raise ValueError(f"wait_signal: unsupported backend: {backend}") +def reduce_scatter_offset( + input: torch.Tensor, + out: list[torch.Tensor], + group: str, + *, + dim: int, + offsets: list[int] | None = None, + dst_ranks: list[int] | None = None, + red_op: str = "sum", +) -> None: + r""" + reduce_scatter_offset(input, out, group, *, dim, offsets, dst_ranks, red_op='sum') -> None + + Simultaneously reduce N blocks of a 2-D ``input`` tensor from a symmetric + memory buffer, routing each block to a specific destination rank. Only + ``dst_ranks[i]`` writes the reduced result for block ``i``; the result is + written to a contiguous output tensor, with the same shape as block ``i``. + + The ``dim`` argument controls which dimension is sharded: + + - ``dim=0`` (row sharding): block ``i`` spans + ``input[offsets[i-1] : offsets[i], :]``. Each ``out[j]`` has shape + ``(size_j, input.size(1))``. + - ``dim=1`` (column sharding): block ``i`` spans + ``input[:, offsets[i-1] : offsets[i]]``. Each ``out[j]`` has shape + ``(input.size(0), size_j)``. + + Blocks are described by ``offsets``, an inclusive prefix-sum of block sizes + along ``dim`` (first block starts at index 0 by convention). Block offsets + can be even or uneven; when uneven, the following condition must be met: for + each ``j``, the ``j``-th owned block must have the same size across all + ranks (so that ``out[j]`` has a uniform shape); different ``j``'s may + differ. + + Args: + input (Tensor): 2-D tensor allocated via symmetric memory (innermost + dimension must be contiguous). + out (list[Tensor]): Output tensors for this rank's owned blocks. Must + have length equal to the number of blocks owned by this rank (i.e. + the count of ``i`` where ``dst_ranks[i] == my_rank``). Each + ``out[j]`` must be contiguous with the same dtype as ``input``. + group (str): The name of the ``ProcessGroup`` to perform the operation on. + dim (int): Dimension along which blocks are defined (0 or 1). + offsets (list[int] | None): Inclusive prefix-sum of block sizes along + ``dim``, length N. If not provided, ``input.size(dim)`` is divided + into equal-size blocks based on the size of the ``group``. + dst_ranks (list[int] | None): Destination rank for each block. If not + provided, blocks are distributed round-robin across ranks. + red_op (str): Reduction operation; currently only ``'sum'`` is supported. + + Example:: + + >>> # doctest: +SKIP + >>> # Each rank holds a Grouped GEMM gradient buffer in symmetric memory. + >>> # The buffer has W experts laid out as equal column blocks; each expert + >>> # is reduced to a specific rank (dst_ranks[i] == i % world_size). + >>> buf = symm_mem.empty(H, W * C, dtype=torch.bfloat16, device="cuda") + >>> symm_mem.rendezvous(buf, group=group_name) + >>> offsets = [i * C for i in range(1, W + 1)] # inclusive prefix-sum + >>> dst_ranks = [i % world_size for i in range(W)] + >>> n_owned = sum(r == rank for r in dst_ranks) + >>> out = [torch.empty(H, C, dtype=torch.bfloat16, device="cuda") for _ in range(n_owned)] + >>> symm_mem.reduce_scatter_offset(buf, out, group_name, dim=1, offsets=offsets, dst_ranks=dst_ranks) + """ + backend = get_backend(input.device) + if backend == "NCCL": + torch.ops.symm_mem.nccl_reduce_scatter_offset( + input, out, group, dim, offsets, dst_ranks, red_op + ) + else: + raise NotImplementedError( + f"reduce_scatter_offset: unsupported backend: {backend}" + ) + + +def is_symm_mem_tensor(tensor: torch.Tensor) -> bool: + r""" + is_symm_mem_tensor(tensor) -> bool + + Returns ``True`` if ``tensor`` was allocated via symmetric memory + (i.e. via :func:`torch.distributed._symmetric_memory.empty` or + :meth:`_SymmetricMemory.empty_strided_p2p`). + + This is a non-collective, O(1) check. + + Args: + tensor (:class:`torch.Tensor`): the tensor to check. + """ + return _SymmetricMemory.is_symm_mem_tensor(tensor) + + __all__ = [ "empty", + "is_symm_mem_tensor", "rendezvous", "is_nvshmem_available", "set_backend", @@ -2204,4 +2298,5 @@ def wait_signal(hdl: _SymmetricMemory, peer: int) -> None: "set_signal_pad_size", "get_signal_pad_size", "get_mem_pool", + "reduce_scatter_offset", ] diff --git a/torch/distributed/_tools/fsdp2_mem_tracker.py b/torch/distributed/_tools/fsdp2_mem_tracker.py index 83d07940a266a..1a10223f03ae4 100644 --- a/torch/distributed/_tools/fsdp2_mem_tracker.py +++ b/torch/distributed/_tools/fsdp2_mem_tracker.py @@ -257,6 +257,7 @@ def inner( state = _FSDPModState.AFT_PRE_FW mod_stat.snapshots.setdefault(state, []).append(self.get_tracker_snapshot()) self._fsdp_state = _FSDPState.FW + # pyrefly: ignore [bad-return] return args, kwargs return inner diff --git a/torch/distributed/checkpoint/filesystem.py b/torch/distributed/checkpoint/filesystem.py index d693790356a9d..831fa16394563 100644 --- a/torch/distributed/checkpoint/filesystem.py +++ b/torch/distributed/checkpoint/filesystem.py @@ -756,7 +756,7 @@ def _write_data( return fut def finish(self, metadata: Metadata, results: list[list[WriteResult]]) -> None: - metadata = dataclasses.replace(metadata, version=CURRENT_DCP_VERSION) + metadata.version = CURRENT_DCP_VERSION storage_md = {} for wr_list in results: diff --git a/torch/distributed/checkpoint/state_dict.py b/torch/distributed/checkpoint/state_dict.py index ed3d4622eace4..7d4428afa25db 100644 --- a/torch/distributed/checkpoint/state_dict.py +++ b/torch/distributed/checkpoint/state_dict.py @@ -844,6 +844,7 @@ def _reconstruct_nested_dict( continue # Reconstruct state for this parameter + # pyrefly: ignore [unsupported-operation] state[fqn] = {} for state_name in optim.state[param]: flattened_state_key = f"{_STATE}.{fqn}.{state_name}" @@ -853,11 +854,13 @@ def _reconstruct_nested_dict( reconstructed_value = _reconstruct_nested_dict( flattened_state_key, state_dict ) + # pyrefly: ignore [bad-index] cast(DictValueType, state[fqn])[state_name] = ( reconstructed_value ) else: # Existing keys mean no nesting, directly use the value. + # pyrefly: ignore [bad-index] cast(DictValueType, state[fqn])[state_name] = state_dict[ flattened_state_key ] diff --git a/torch/distributed/checkpoint/state_dict_loader.py b/torch/distributed/checkpoint/state_dict_loader.py index 62acd1bfef359..955ccbcdcf3bc 100644 --- a/torch/distributed/checkpoint/state_dict_loader.py +++ b/torch/distributed/checkpoint/state_dict_loader.py @@ -228,11 +228,15 @@ def local_step(): nonlocal metadata # Use global metadata if available, otherwise fallback to rank local metadata + global_metadata_exc: Exception | None = None + rank_metadata_exc: Exception | None = None try: metadata = storage_reader.read_metadata() - except Exception: - logger.info( - "Global metadata is not found. Falling back to rank local metadata." + except Exception as e: + global_metadata_exc = e + logger.warning( + "Global metadata is not found. Falling back to rank local metadata.", + exc_info=True, ) if ( @@ -242,13 +246,23 @@ def local_step(): try: metadata = storage_reader.read_metadata(rank=distW.rank) # noqa: F841 use_collectives = False - except Exception: - logger.info("Rank local metadata is not found.") + except Exception as e: + rank_metadata_exc = e + logger.warning("Rank local metadata is not found.", exc_info=True) if planner is None: raise AssertionError("planner is None") if metadata is None: - raise AssertionError("metadata is None") + error_parts = ["metadata is None"] + if global_metadata_exc is not None: + error_parts.append( + f"global metadata read failed: {global_metadata_exc}" + ) + if rank_metadata_exc is not None: + error_parts.append( + f"rank local metadata read failed: {rank_metadata_exc}" + ) + raise AssertionError("; ".join(error_parts)) planner.set_up_planner(state_dict, metadata, distW.is_coordinator) if ( diff --git a/torch/distributed/checkpoint/utils.py b/torch/distributed/checkpoint/utils.py index d5e19eee111a0..4f14bce83b09e 100644 --- a/torch/distributed/checkpoint/utils.py +++ b/torch/distributed/checkpoint/utils.py @@ -261,6 +261,7 @@ def all_reduce( final_result = self.broadcast_object(result) if isinstance(final_result, CheckpointException): raise final_result + # pyrefly: ignore [redundant-cast] return cast(R, final_result) def all_gather( @@ -310,6 +311,7 @@ def broadcast( final_result = self.broadcast_object(result) if isinstance(final_result, CheckpointException): raise final_result + # pyrefly: ignore [redundant-cast] return cast(T, final_result) def barrier(self) -> None: diff --git a/torch/distributed/debug/_debug_handlers.py b/torch/distributed/debug/_debug_handlers.py index f63c0848d7e15..902f78e2e863f 100644 --- a/torch/distributed/debug/_debug_handlers.py +++ b/torch/distributed/debug/_debug_handlers.py @@ -228,9 +228,10 @@ def templates(self) -> dict[str, str]: return {"pyspy_dump.html": PYSPY_DUMP_TEMPLATE} def _handle(self, req: HTTPRequestHandler) -> bytes: - addrs, resps = fetch_all( - "pyspy_dump", req.get_raw_query(), timeout=self.fetch_timeout - ) + query = req.get_raw_query() + if "nonblocking" not in query: + query = f"nonblocking=1&{query}" if query else "nonblocking=1" + addrs, resps = fetch_all("pyspy_dump", query, timeout=self.fetch_timeout) return req.frontend.render_template( "pyspy_dump.html", addrs=addrs, @@ -238,7 +239,9 @@ def _handle(self, req: HTTPRequestHandler) -> bytes: ) def dump(self) -> str | None: - addrs, resps = fetch_all("pyspy_dump", timeout=self.fetch_timeout) + addrs, resps = fetch_all( + "pyspy_dump", "nonblocking=1", timeout=self.fetch_timeout + ) parts: list[str] = [] summary = format_fetch_summary(addrs, resps) if summary: diff --git a/torch/distributed/device_mesh.py b/torch/distributed/device_mesh.py index c84544d2cceaa..c0541896b094d 100644 --- a/torch/distributed/device_mesh.py +++ b/torch/distributed/device_mesh.py @@ -4,7 +4,7 @@ import os import threading import warnings -from collections.abc import Iterator +from collections.abc import Callable, Iterator from itertools import zip_longest from typing import Any, TYPE_CHECKING @@ -549,10 +549,12 @@ def _init_one_process_group( getattr(default_group, "bound_device_id", None) is not None or dist_config.use_torchcomms ) - and torch.cuda.is_available() + and torch.accelerator.is_available() and ( backend is None - or default_group._get_backend(torch.device("cuda")).name() + or default_group._get_backend( + torch.accelerator.current_accelerator() # pyrefly: ignore[bad-argument-type] + ).name() == backend ) ): @@ -826,41 +828,49 @@ def _create_sub_mesh( layout: _MeshLayout, submesh_dim_names: tuple[str, ...], ) -> "DeviceMesh": - root_mesh = self._get_root_mesh() - slice_dim_group_name = [] - if len(self._dim_group_names) > 0: - if len(self._dim_group_names) != len(not_none(self._mesh_dim_names)): - raise AssertionError( - "The number of dim_group_names and mesh_dim_names " - "should have the same length if the rank is in the mesh." - ) - for name in submesh_dim_names: - if name in not_none(self._mesh_dim_names): - slice_dim_group_name.append( - self._dim_group_names[ - not_none(self._mesh_dim_names).index(name) - ] - ) - else: - # If device_mesh is not root_mesh, we already throw error in _get_slice_mesh_layout - # Since we will deprecate the slicing of flattened dim_name from root mesh soon, - # we don't want to optimize the code furthermore. - flatten_mesh = self._flatten_mapping[name] - slice_dim_group_name.append( - flatten_mesh._dim_group_names[ - not_none(flatten_mesh._mesh_dim_names).index(name) - ] + # DeviceMesh() always graph breaks (its __init__ is skipped by + # dynamo). With nested graph breaks the subsequent STORE_ATTR on + # the new opaque object ends up in a resume function and triggers a + # hard error. Disable nested graph breaks so the break propagates + # to the caller instead. + with torch._dynamo.disable_nested_graph_breaks(): + root_mesh = self._get_root_mesh() + slice_dim_group_name = [] + if len(self._dim_group_names) > 0: + if len(self._dim_group_names) != len( + not_none(self._mesh_dim_names) + ): + raise AssertionError( + "The number of dim_group_names and mesh_dim_names " + "should have the same length if the rank is in the mesh." ) - res_submesh = DeviceMesh( - self._device_type, - _layout=layout, - _rank_map=root_mesh._rank_map, - mesh_dim_names=submesh_dim_names, - _root_mesh=root_mesh, - _init_backend=False, - ) - res_submesh._dim_group_names = slice_dim_group_name - return res_submesh + for name in submesh_dim_names: + if name in not_none(self._mesh_dim_names): + slice_dim_group_name.append( + self._dim_group_names[ + not_none(self._mesh_dim_names).index(name) + ] + ) + else: + # If device_mesh is not root_mesh, we already throw error in _get_slice_mesh_layout + # Since we will deprecate the slicing of flattened dim_name from root mesh soon, + # we don't want to optimize the code furthermore. + flatten_mesh = self._flatten_mapping[name] + slice_dim_group_name.append( + flatten_mesh._dim_group_names[ + not_none(flatten_mesh._mesh_dim_names).index(name) + ] + ) + res_submesh = DeviceMesh( + self._device_type, + _layout=layout, + _rank_map=root_mesh._rank_map, + mesh_dim_names=submesh_dim_names, + _root_mesh=root_mesh, + _init_backend=False, + ) + res_submesh._dim_group_names = slice_dim_group_name + return res_submesh def _create_flatten_mesh( self, @@ -1240,7 +1250,11 @@ def _sym_get_coordinate(self, index: int) -> IntLikeType: import torch.distributed.config as config from torch._guards import detect_fake_mode - if not detect_fake_mode() or not config.compile_on_one_rank: + if ( + not config.compile_on_one_rank + or not (fake_mode := detect_fake_mode()) + or not fake_mode.shape_env + ): # This is only valid when the current rank is part of the mesh. if self._coordinate_on_dim is None: raise AssertionError @@ -1590,12 +1604,89 @@ def init_device_mesh( return device_mesh +_distributed_opaque_types_registered = False + + +def _device_mesh_reconstruct_fn( + mesh: "OpaqueBase", + get_tracked_proxy: Callable[["OpaqueBase"], "torch.fx.Proxy | None"], + tracer: Any, +) -> "torch.fx.Proxy | None": + """Reconstruct a DeviceMesh submesh from a tracked ancestor mesh. + + Called by PythonKeyTracer when make_fx encounters a DeviceMesh that isn't + tracked (e.g. a submesh captured by a backward closure). Looks for any + tracked mesh that shares the same root and contains the target dim names, + then emits a call_function node that derives the submesh via _get_submesh. + """ + if not isinstance(mesh, DeviceMesh): + raise AssertionError("DeviceMesh expected") + + root_mesh = mesh._get_root_mesh() + + # Only submeshes can be reconstructed; root meshes must already be tracked. + if mesh is root_mesh: + return None + + dim_names = mesh._mesh_dim_names + if dim_names is None: + return None + + # Ensure the custom ops are registered + from torch.distributed._ops import device_mesh as _dm_ops # noqa: F401 + + # Try the root mesh first (original path). + ancestor_proxy = get_tracked_proxy(root_mesh) + ancestor_dim_names = root_mesh._mesh_dim_names + + # If root isn't tracked, search for any tracked DeviceMesh that shares + # the same root AND contains all our dim names. This handles the case + # where e.g. a concatenated (fsdp, tp) mesh is a graph input (from + # DTensor.__tensor_flatten__) but neither root nor the individual + # submeshes are tracked directly. + if ancestor_proxy is None: + from torch._library.fake_class_registry import FakeScriptObject + + for tracked_obj, proxy in tracer.opaque_tracker.items(): + real_obj = ( + tracked_obj.real_obj + if isinstance(tracked_obj, FakeScriptObject) + else tracked_obj + ) + if not isinstance(real_obj, DeviceMesh) or real_obj is mesh: + continue + if real_obj._get_root_mesh() is not root_mesh: + continue + tracked_dim_names = real_obj._mesh_dim_names + if tracked_dim_names is None: + continue + if all(n in tracked_dim_names for n in dim_names): + ancestor_proxy = proxy + ancestor_dim_names = tracked_dim_names + break + + if ancestor_proxy is None or ancestor_dim_names is None: + return None + + # Convert our dim names to indices into the ancestor mesh's dim names + mesh_dims = [ancestor_dim_names.index(n) for n in dim_names] + + # Dispatch through the custom op with proxy mode active so that + # meta["val"] is set and the result is tracked in opaque_tracker. + return torch.ops.device_mesh._get_submesh(ancestor_proxy, mesh_dims) + + def _register_distributed_opaque_types(): """ Register DeviceMesh as an opaque type for torch.compile. This must happen before any custom ops that use DeviceMesh in their schema. Called lazily to avoid circular import issues. """ + global _distributed_opaque_types_registered + if _distributed_opaque_types_registered: + return + _distributed_opaque_types_registered = True + from torch._library.opaque_object import MemberType, register_opaque_type register_opaque_type( @@ -1608,12 +1699,14 @@ def _register_distributed_opaque_types(): "group_name": MemberType.USE_REAL, "group_desc": MemberType.USE_REAL, "__eq__": MemberType.USE_REAL, + "__ne__": MemberType.USE_REAL, }, ) register_opaque_type( DeviceMesh, typ="reference", + reconstruct_fn=_device_mesh_reconstruct_fn, guard_fn=lambda obj: [ obj._flatten_rank_map, obj._layout, @@ -1634,6 +1727,7 @@ def _register_distributed_opaque_types(): "get_coordinate": MemberType.USE_REAL, "get_local_rank": MemberType.USE_REAL, "__eq__": MemberType.USE_REAL, + "__ne__": MemberType.USE_REAL, "ndim": MemberType.USE_REAL, "shape": MemberType.USE_REAL, "mesh_dim_names": MemberType.USE_REAL, diff --git a/torch/distributed/distributed_c10d.py b/torch/distributed/distributed_c10d.py index 42579f13bcbb0..05986c191a233 100644 --- a/torch/distributed/distributed_c10d.py +++ b/torch/distributed/distributed_c10d.py @@ -134,6 +134,7 @@ "get_node_local_rank", "split_group", "shrink_group", + "record_comm", ] _MPI_AVAILABLE = True @@ -143,11 +144,15 @@ _XCCL_AVAILABLE = True try: - # pyrefly: ignore [missing-import] - from torchcomms._backend_wrapper import _BackendWrapper + try: + # pyrefly: ignore [missing-import] + from torchcomms._comms import _BackendWrapper + except ImportError: + # pyrefly: ignore [missing-import] + from torchcomms._backend_wrapper import _BackendWrapper # pyrefly: ignore [missing-import] - from torchcomms._comms import new_comm + from torchcomms import new_comm # pyrefly: ignore [missing-import] from torchcomms.hooks import FlightRecorderHook @@ -5161,6 +5166,7 @@ def barrier( group: ProcessGroup | None = GroupMember.WORLD, async_op: bool = False, device_ids=None, + timeout: timedelta | None = None, ): """ Synchronize all processes. @@ -5173,6 +5179,8 @@ def barrier( the default process group will be used. async_op (bool, optional): Whether this op should be an async op device_ids ([int], optional): List of device/GPU ids. Only one id is expected. + timeout (datetime.timedelta, optional): Timeout for barrier. + If ``None``, the default process group timeout will be used. Returns: Async work handle, if async_op is set to True. @@ -5193,6 +5201,8 @@ def barrier( opts = BarrierOptions() opts.asyncOp = async_op + if timeout is not None: + opts.timeout = timeout # Detect the accelerator on the machine. If no accelerator is available, it # returns CPU. device = torch._C._get_accelerator() @@ -5453,10 +5463,18 @@ def split_group( ) parent_group_rank = parent_global_to_group_ranks[global_rank] - parent_backend = parent_pg._get_backend(torch.device("cuda")) + + if torch.accelerator.is_available(): + parent_backend = parent_pg._get_backend( + torch.accelerator.current_accelerator() # pyrefly: ignore[bad-argument-type] + ) + else: + raise RuntimeError( + "No backend for the parent process group or its backend does not support splitting" + ) # if the parent backend does not support splitting, raise error - # currently this API only support NCCL backend + # currently this API only support NCCL and XCCL backend if ( not parent_backend or not parent_backend.supports_splitting ) and not _use_torchcomms_enabled(): @@ -5522,7 +5540,16 @@ def split_group( global_ranks_in_my_group = [parent_group_to_global_ranks[rank] for rank in my_group] split_pg.bound_device_id = device_id # type: ignore[union-attr] - split_backend_class = split_pg._get_backend(torch.device("cuda")) + + if torch.accelerator.is_available(): + split_backend_class = split_pg._get_backend( + torch.accelerator.current_accelerator() # pyrefly: ignore[bad-argument-type] + ) + else: + raise RuntimeError( + "No backend for the parent process group or its backend does not support splitting" + ) + if not _use_torchcomms_enabled(): split_backend_class._set_sequence_number_for_group() if split_pg.group_name != group_name: @@ -6555,3 +6582,28 @@ def _update_process_group_global_state( # Standard process group tag _world.tags_to_pg.setdefault(pg_tag, []).append(pg) _world.pg_to_tag[pg] = pg_tag + + +@contextlib.contextmanager +def record_comm(name: str): + """Context manager to set a custom profiling name for communication collectives. + + When active, all c10d collectives issued within this context will use ``name`` + as their profiling title in the Work base class, overriding the default + backend-specific name (e.g. ``nccl:all_reduce``). This works across all + backends without per-backend or per-collective changes. + + Args: + name (str): The profiling name to associate with collectives. + + Example:: + >>> # xdoctest: +SKIP("undefined vars") + >>> with dist.record_comm("FSDP::all_gather (layer1)"): + ... dist.all_gather_into_tensor(output, input, group=pg) + """ + prev = torch._C._distributed_c10d._get_comm_profiling_name() + torch._C._distributed_c10d._set_comm_profiling_name(name) + try: + yield + finally: + torch._C._distributed_c10d._set_comm_profiling_name(prev) diff --git a/torch/distributed/elastic/agent/server/api.py b/torch/distributed/elastic/agent/server/api.py index dd0905eb3f761..fe3dda0841a1d 100644 --- a/torch/distributed/elastic/agent/server/api.py +++ b/torch/distributed/elastic/agent/server/api.py @@ -471,6 +471,7 @@ def __init__( self._exit_barrier_timeout = exit_barrier_timeout self._shutdown_timeout = shutdown_timeout self._total_execution_time = 0 + self._in_exit_barrier: bool = False def get_worker_group(self, role: str = DEFAULT_ROLE) -> WorkerGroup: return self._worker_group @@ -902,19 +903,9 @@ def _record_flakiness_metric(self, is_failed: bool = False): put_metric(f"workers.{spec.role}.flakiness", int(flakiness)) - def _pre_invoke_run(self) -> None: - """Hook called before the worker lifecycle loop in ``_invoke_run``. - - Subclasses can override this to perform setup that must happen - before rendezvous and worker initialization (e.g. starting a - health check server). The default implementation is a no-op. - """ - def _invoke_run(self, role: str = DEFAULT_ROLE) -> RunResult: # NOTE: currently only works for a single role - self._pre_invoke_run() - spec = self._worker_group.spec role = spec.role @@ -998,6 +989,7 @@ def _exit_barrier(self): self._exit_barrier_timeout, ) start = time.time() + self._in_exit_barrier = True try: store_util.barrier( store=self._store, @@ -1017,3 +1009,5 @@ def _exit_barrier(self): "Error waiting on exit barrier. Elapsed: %s seconds", time.time() - start, ) + finally: + self._in_exit_barrier = False diff --git a/torch/distributed/elastic/agent/server/local_elastic_agent.py b/torch/distributed/elastic/agent/server/local_elastic_agent.py index b9d2ed64cd0cf..fda632bb3aaba 100644 --- a/torch/distributed/elastic/agent/server/local_elastic_agent.py +++ b/torch/distributed/elastic/agent/server/local_elastic_agent.py @@ -7,6 +7,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from __future__ import annotations import json import os @@ -43,6 +44,8 @@ if TYPE_CHECKING: + from collections.abc import Callable + from torch.distributed.elastic.events.api import EventMetadataValue logger = get_logger(__name__) @@ -59,6 +62,29 @@ TORCHELASTIC_TIMER_FILE = "TORCHELASTIC_TIMER_FILE" +class _AliveCallbackProxy: + """Mutable callback wrapper for the health check server. + + The C++ pybind ``HealthCheckThriftServer`` binds its ``alive_callback`` + at construction time and cannot update it afterward. This proxy is + created *before* the health check server so it can be passed as the + callback. Initially it returns ``time.time()`` (signalling "alive"). + After the agent is constructed, :meth:`set_delegate` wires it to + ``agent._get_alive_time`` for real liveness tracking. + """ + + def __init__(self) -> None: + self._delegate: Callable[[], int] | None = None + + def __call__(self) -> int: + if self._delegate is not None: + return self._delegate() + return int(time.time()) + + def set_delegate(self, delegate: Callable[[], int]) -> None: + self._delegate = delegate + + class LocalElasticAgent(SimpleElasticAgent): """An implementation of :py:class:`torchelastic.agent.server.ElasticAgent` that handles host-local workers. @@ -156,6 +182,7 @@ def __init__( exit_barrier_timeout: float = 300, log_line_prefix_template: str | None = None, shutdown_timeout: int = 30, + health_check_server: HealthCheckServer | None = None, ): super().__init__(spec, exit_barrier_timeout, shutdown_timeout) self._start_method = start_method @@ -164,7 +191,7 @@ def __init__( self._log_line_prefix_template = log_line_prefix_template self._worker_watchdog: timer.FileTimerServer | None = None self._logs_specs = logs_specs - self._health_check_server: HealthCheckServer | None = None + self._health_check_server = health_check_server def _setup_local_watchdog(self, envs: dict[int, dict[str, str]]) -> None: enable_watchdog_env_name = TORCHELASTIC_ENABLE_FILE_TIMER @@ -217,7 +244,13 @@ def _get_alive_time(self) -> int: Once workers are running and the watchdog is active, we delegate to the watchdog's ``get_last_progress_time`` for real liveness tracking. + + During the exit barrier wait, workers have finished and the watchdog + progress time is stale. We return the current time to prevent TW + from killing the task while agents coordinate shutdown. """ + if self._in_exit_barrier: + return int(time.time()) if self._worker_watchdog is not None: return self._worker_watchdog.get_last_progress_time() return int(time.time()) @@ -440,22 +473,6 @@ def _set_local_rank_env( if "CUDA_VISIBLE_DEVICES" in os.environ: worker_env["CUDA_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"] - def _pre_invoke_run(self) -> None: - # Start the health check server immediately, before rendezvous, - # so that TW sees a healthy thrift port while the agent is still - # initializing (package fetch, rendezvous, model load, etc.). - # The _get_alive_time callback dynamically checks self._worker_watchdog: - # returns current time during init, real watchdog time once workers run. - if justknobs_check( - "ai_infra/pytorch_distributed:torchelastic_enable_healthcheck_before_rendezvous", - default=False, - ): - logger.info( - "Starting health check server before rendezvous " - "(torchelastic_enable_healthcheck_before_rendezvous=True)" - ) - self._setup_healthcheck() - def _shutdown( self, death_sig: signal.Signals = signal.SIGTERM, timeout: int = 30 ) -> None: diff --git a/torch/distributed/elastic/multiprocessing/redirects.py b/torch/distributed/elastic/multiprocessing/redirects.py index 057013fbb9e5b..cecdec12ffee2 100644 --- a/torch/distributed/elastic/multiprocessing/redirects.py +++ b/torch/distributed/elastic/multiprocessing/redirects.py @@ -24,12 +24,28 @@ logger = logging.getLogger(__name__) +_WIN32_STD_HANDLE = { + "stdout": -11, # STD_OUTPUT_HANDLE + "stderr": -12, # STD_ERROR_HANDLE +} + + def get_libc(): - if IS_WINDOWS or IS_MACOS: - logger.warning( - "NOTE: Redirects are currently not supported in Windows or MacOs." - ) + if IS_MACOS: + logger.warning("NOTE: Redirects are currently not supported in MacOs.") return None + elif IS_WINDOWS: + for lib_name in ("ucrtbase", "msvcrt", "msvcr110", "msvcr100"): + try: + lib = ctypes.CDLL(lib_name) + logger.debug("Loaded Windows C runtime: %s", lib_name) + return lib + except OSError: + continue + raise RuntimeError( + "Could not load a C runtime DLL on Windows (tried: ucrtbase, msvcrt, " + "msvcr110, msvcr100). Redirects cannot function without a CRT." + ) else: return ctypes.CDLL("libc.so.6") @@ -38,6 +54,24 @@ def get_libc(): def _c_std(stream: str): + if IS_WINDOWS: + stream_index = 2 if stream == "stderr" else 1 + try: + iob_func = libc.__acrt_iob_func + iob_func.restype = ctypes.POINTER(ctypes.c_void_p) + iob_func.argtypes = [ctypes.c_uint] + return iob_func(stream_index) + except AttributeError: + pass + try: + legacy_index = 2 if stream == "stderr" else 1 + iob = (ctypes.POINTER(ctypes.c_void_p) * 3).in_dll(libc, "_iob") + return iob[legacy_index] + except (AttributeError, OSError) as err: + raise RuntimeError( + f"Could not resolve C-runtime FILE* for '{stream}'. " + "Neither __acrt_iob_func nor _iob are available in the loaded CRT." + ) from err return ctypes.c_void_p.in_dll(libc, stream) @@ -48,56 +82,147 @@ def _python_std(stream: str): _VALID_STD = {"stdout", "stderr"} -@contextmanager -def redirect(std: str, to_file: str): - """ - Redirect ``std`` (one of ``"stdout"`` or ``"stderr"``) to a file in the path specified by ``to_file``. - - This method redirects the underlying std file descriptor (not just python's ``sys.stdout|stderr``). - See usage for details. - - Directory of ``dst_filename`` is assumed to exist and the destination file - is overwritten if it already exists. - - .. note:: Due to buffering cross source writes are not guaranteed to - appear in wall-clock order. For instance in the example below - it is possible for the C-outputs to appear before the python - outputs in the log file. - - Usage: - - :: - - # syntactic-sugar for redirect("stdout", "tmp/stdout.log") - with redirect_stdout("/tmp/stdout.log"): - print("python stdouts are redirected") - libc = ctypes.CDLL("libc.so.6") - libc.printf(b"c stdouts are also redirected" - os.system("echo system stdouts are also redirected") - - print("stdout restored") - - """ - if std not in _VALID_STD: - raise ValueError( - f"unknown standard stream <{std}>, must be one of {_VALID_STD}" - ) - - c_std = _c_std(std) - python_std = _python_std(std) - std_fd = python_std.fileno() - - def _redirect(dst): - libc.fflush(c_std) - python_std.flush() - os.dup2(dst.fileno(), std_fd) - - with os.fdopen(os.dup(std_fd)) as orig_std, open(to_file, mode="w+b") as dst: - _redirect(dst) - try: - yield - finally: - _redirect(orig_std) +if IS_WINDOWS: # libc is None on macOS; all of the below is Windows-only + import io as _io + import msvcrt as _msvcrt + + _kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) # type: ignore[attr-defined] + + _crt_dup = libc._dup + _crt_dup2 = libc._dup2 + _crt_dup.restype = ctypes.c_int + _crt_dup.argtypes = [ctypes.c_int] + _crt_dup2.restype = ctypes.c_int + _crt_dup2.argtypes = [ctypes.c_int, ctypes.c_int] + + @contextmanager + def redirect(std: str, to_file: str): + """ + Redirect ``std`` (one of ``"stdout"`` or ``"stderr"``) to a file at ``to_file``. + + On Windows this performs a four-layer redirect: + + 1. ``sys.stdout``/``sys.stderr`` -- rewired to a new TextIOWrapper so + Python's ``print()`` writes to the destination file. + 2. CRT fd (``_dup2``) -- captures C ``printf`` and UCRT ``FILE*`` writers. + 3. Win32 ``SetStdHandle`` -- captures native code using ``WriteFile``/ + ``WriteConsole`` directly, including HIP/ROCm. + 4. ``fflush`` before each switch -- prevents lost output from CRT buffering. + + .. note:: If ROCm/HIP caches the Win32 HANDLE before this redirect runs + (e.g. at ``import torch`` time), set up the redirect *before* + importing torch/ROCm to capture all output. + + Directory of ``to_file`` is assumed to exist. The destination file is + overwritten if it already exists. + """ + if std not in _VALID_STD: + raise ValueError( + f"unknown standard stream <{std}>, must be one of {_VALID_STD}" + ) + + std_fd = 1 if std == "stdout" else 2 + win32_handle_id = _WIN32_STD_HANDLE[std] + orig_sys_std = getattr(sys, std) + orig_fd_dup = _crt_dup(std_fd) + if orig_fd_dup == -1: + raise OSError(f"CRT _dup failed for {std} (fd={std_fd})") + orig_win32_handle = _kernel32.GetStdHandle(win32_handle_id) + + with open(to_file, mode="w+b") as dst: + dst_fd = dst.fileno() + + try: + libc.fflush(_c_std(std)) + except Exception: + pass + try: + orig_sys_std.flush() + except Exception: + pass + + _kernel32.SetStdHandle( + win32_handle_id, + _msvcrt.get_osfhandle(dst_fd), # pyrefly: ignore [missing-attribute] + ) + + if _crt_dup2(dst_fd, std_fd) == -1: + raise OSError(f"CRT _dup2 failed redirecting {std}") + + new_sys_std = _io.TextIOWrapper( + open(dst_fd, mode="wb", closefd=False), # noqa: SIM115 + encoding=orig_sys_std.encoding or "utf-8", + errors="replace", + line_buffering=True, + ) + setattr(sys, std, new_sys_std) + + try: + yield + finally: + try: + new_sys_std.flush() + except Exception: + pass + try: + libc.fflush(_c_std(std)) + except Exception: + pass + + setattr(sys, std, orig_sys_std) + _crt_dup2(orig_fd_dup, std_fd) + os.close(orig_fd_dup) + _kernel32.SetStdHandle(win32_handle_id, orig_win32_handle) + +else: + + @contextmanager + def redirect(std: str, to_file: str): + """ + Redirect ``std`` (one of ``"stdout"`` or ``"stderr"``) to a file in the path specified by ``to_file``. + + This method redirects the underlying std file descriptor (not just python's ``sys.stdout|stderr``). + See usage for details. + + Directory of ``dst_filename`` is assumed to exist and the destination file + is overwritten if it already exists. + + .. note:: Due to buffering cross source writes are not guaranteed to + appear in wall-clock order. For instance in the example below + it is possible for the C-outputs to appear before the python + outputs in the log file. + + Usage:: + + # syntactic-sugar for redirect("stdout", "tmp/stdout.log") + with redirect_stdout("/tmp/stdout.log"): + print("python stdouts are redirected") + libc = ctypes.CDLL("libc.so.6") + libc.printf(b"c stdouts are also redirected") + os.system("echo system stdouts are also redirected") + + print("stdout restored") + """ + if std not in _VALID_STD: + raise ValueError( + f"unknown standard stream <{std}>, must be one of {_VALID_STD}" + ) + + c_std = _c_std(std) + python_std = _python_std(std) + std_fd = python_std.fileno() + + def _redirect(dst): + libc.fflush(c_std) + python_std.flush() + os.dup2(dst.fileno(), std_fd) + + with os.fdopen(os.dup(std_fd)) as orig_std, open(to_file, mode="w+b") as dst: + _redirect(dst) + try: + yield + finally: + _redirect(orig_std) redirect_stdout = partial(redirect, "stdout") diff --git a/torch/distributed/elastic/rendezvous/etcd_server.py b/torch/distributed/elastic/rendezvous/etcd_server.py index 347e7339d9a46..c3a6b2803d85d 100644 --- a/torch/distributed/elastic/rendezvous/etcd_server.py +++ b/torch/distributed/elastic/rendezvous/etcd_server.py @@ -177,7 +177,7 @@ def start( curr_retries += 1 stop_etcd(self._etcd_proc) logger.warning( # noqa: G200 - "Failed to start etcd server, got error: %s, retrying", str(e) + "Failed to start etcd server, got error: %s, retrying", e ) if curr_retries >= num_retries: shutil.rmtree(self._base_data_dir, ignore_errors=True) diff --git a/torch/distributed/elastic/rendezvous/etcd_store.py b/torch/distributed/elastic/rendezvous/etcd_store.py index faaf77587bc9d..09e588ff5b709 100644 --- a/torch/distributed/elastic/rendezvous/etcd_store.py +++ b/torch/distributed/elastic/rendezvous/etcd_store.py @@ -120,6 +120,7 @@ def add(self, key, num: int) -> int: except etcd.EtcdCompareFailed: cas_delay() + # pyrefly: ignore [bad-override] def wait(self, keys, override_timeout: datetime.timedelta | None = None): """ Wait until all of the keys are published, or until timeout. diff --git a/torch/distributed/flight_recorder/components/builder.py b/torch/distributed/flight_recorder/components/builder.py index 87c52d6d4b090..663363e718257 100644 --- a/torch/distributed/flight_recorder/components/builder.py +++ b/torch/distributed/flight_recorder/components/builder.py @@ -311,7 +311,12 @@ def build_collectives( # This extra cleanup is needed because we need to pop all collectives within a coalesced collective. for i, k in idx_map.items(): for _ in range(1, num_coalesced_entries): - all_entries[i].pop(k) + try: + all_entries[i].pop(k) + except IndexError: + # In the case of a missing rank symptom that a rank didn't schedule the coalesced collective, + # we should not fail the analysis script here. + pass else: # Iterate through all the ranks and check if there is a mismatch for the current entry. check_current_entry_match( diff --git a/torch/distributed/flight_recorder/components/types.py b/torch/distributed/flight_recorder/components/types.py index d0b8f067acfb4..46ebda50dd506 100644 --- a/torch/distributed/flight_recorder/components/types.py +++ b/torch/distributed/flight_recorder/components/types.py @@ -209,16 +209,25 @@ class Database(NamedTuple): "reduce", "_reduce_oop", "all_gather", + "all_gather_single", + "all_gather_v", "all_reduce", "_all_gather_base", "all_gather_into_tensor_coalesced", "reduce_scatter", + "reduce_scatter_single", + "reduce_scatter_v", "reduce_scatter_tensor_coalesced", "_reduce_scatter_base", "gather", "scatter", "all_to_all", + "all_to_all_single", + "all_to_all_v_single", "all_reduce_barrier", + "barrier", + "split", + "new_window", "allreduce_coalesced", "ALLGATHER_coalesced", "REDUCE_SCATTER_coalesced", @@ -283,7 +292,7 @@ def log( logger.info("input sizes: %s", self.input_sizes) logger.info("output sizes: %s", self.output_sizes) logger.info("world size: %d", len(self.expected_ranks)) - logger.info("expected ranks: %s", str(self.expected_ranks)) + logger.info("expected ranks: %s", self.expected_ranks) logger.info("collective state: %s", self.collective_state) if errors: self.errors = errors @@ -409,9 +418,9 @@ def __init__( ): self.profiling_name = event["profiling_name"] comm_lib_backend, name = self.profiling_name.split(":") - if comm_lib_backend not in ["nccl", "xccl"]: + if comm_lib_backend not in ["nccl", "ncclx", "gloo", "xccl"]: raise AssertionError( - f"name formatting error? {comm_lib_backend} != 'nccl' or 'xccl'" + f"name formatting error? {comm_lib_backend} not in supported backends" ) parts = name.split(" ") type = parts[0] diff --git a/torch/distributed/fsdp/_flat_param.py b/torch/distributed/fsdp/_flat_param.py index 190e2a0a88a34..c7e226c09457f 100644 --- a/torch/distributed/fsdp/_flat_param.py +++ b/torch/distributed/fsdp/_flat_param.py @@ -29,6 +29,7 @@ ) from torch.nn.parameter import _ParameterMeta # type: ignore[attr-defined] from torch.testing._internal.distributed.fake_pg import FakeProcessGroup +from torch.utils._typing_utils import not_none from ._fsdp_extensions import ( _ext_post_unflatten_transform, @@ -600,6 +601,7 @@ def __init__( self._needs_pre_backward_unshard = False # Was the handle prefetched? Set on successful _prefetch_handle and unshard self._prefetched = False + self._compute_stream: torch.Stream | None = None # Optimistically assume a valid input `params` and set dtype attributes # before `_init_flat_param()`, which performs the actual validation self._orig_param_dtype = params[0].dtype @@ -1322,6 +1324,11 @@ def pre_unshard(self) -> bool: self._use_sharded_views() ret = False if self._use_orig_params and not self._skip_writeback_check: + # Wait for the compute stream since _writeback_orig_params reads + # original parameters that may still be in use during prefetch. + self._device_handle.current_stream().wait_stream( + not_none(self._compute_stream) + ) ret = self._writeback_orig_params() if ( self.uses_sharded_strategy diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py b/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py index 3c8aa312c7187..f001de1e92d3e 100644 --- a/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py +++ b/torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py @@ -20,6 +20,12 @@ from ._fsdp_param import FSDPParam, ShardedState +def _label_with_suffix(label: str, suffix: str) -> str: + if suffix: + return f"{label} {suffix}" + return label + + class AllGatherResult(NamedTuple): all_gather_output: torch.Tensor all_gather_event: torch.Event | None @@ -330,43 +336,49 @@ def foreach_all_gather( all_gather_stream: torch.Stream, device: torch.device, all_gather_comm: AllGather, + label_suffix: str = "", ) -> AllGatherResult | None: world_size, rank = group.size(), group.rank() device_handle = _get_device_handle(device.type) + with device_handle.stream(all_gather_copy_in_stream): - param_all_gather_inputs = _get_param_all_gather_inputs(fsdp_params) - ( - param_all_gather_input_dtypes, - param_all_gather_input_numels, - dtype, - ) = _get_all_gather_input_metadatas(param_all_gather_inputs) - if dtype == torch.uint8: - all_gather_inputs = [ - t.view(torch.uint8) for ts in param_all_gather_inputs for t in ts - ] - else: - all_gather_inputs = [*chain.from_iterable(param_all_gather_inputs)] - inp_split_sizes = [t.numel() for t in all_gather_inputs] - all_gather_input_numel = sum(inp_split_sizes) - all_gather_output = all_gather_comm.allocate( - (all_gather_input_numel * world_size,), dtype=dtype, device=device - ) - all_gather_input, all_gather_output = torch.ops.fsdp.all_gather_copy_in( - all_gather_inputs, - all_gather_output, - inp_split_sizes, - all_gather_input_numel, - rank, - ) - del param_all_gather_inputs + with torch.profiler.record_function( + _label_with_suffix("FSDP::all_gather_copy_in", label_suffix) + ): + param_all_gather_inputs = _get_param_all_gather_inputs(fsdp_params) + ( + param_all_gather_input_dtypes, + param_all_gather_input_numels, + dtype, + ) = _get_all_gather_input_metadatas(param_all_gather_inputs) + if dtype == torch.uint8: + all_gather_inputs = [ + t.view(torch.uint8) for ts in param_all_gather_inputs for t in ts + ] + else: + all_gather_inputs = [*chain.from_iterable(param_all_gather_inputs)] + inp_split_sizes = [t.numel() for t in all_gather_inputs] + all_gather_input_numel = sum(inp_split_sizes) + all_gather_output = all_gather_comm.allocate( + (all_gather_input_numel * world_size,), dtype=dtype, device=device + ) + all_gather_input, all_gather_output = torch.ops.fsdp.all_gather_copy_in( + all_gather_inputs, + all_gather_output, + inp_split_sizes, + all_gather_input_numel, + rank, + ) + del param_all_gather_inputs all_gather_stream.wait_stream(all_gather_copy_in_stream) with device_handle.stream(all_gather_stream): - all_gather_work = all_gather_comm( - output_tensor=all_gather_output, - input_tensor=all_gather_input, - group=group, - async_op=async_op, - ) + with dist.record_comm(_label_with_suffix("FSDP::all_gather", label_suffix)): + all_gather_work = all_gather_comm( + output_tensor=all_gather_output, + input_tensor=all_gather_input, + group=group, + async_op=async_op, + ) all_gather_event = all_gather_stream.record_event() return AllGatherResult( all_gather_output, @@ -537,6 +549,7 @@ def foreach_reduce( partial_reduce_output: torch.Tensor | None, # only used for HSDP all_reduce_hook: Callable[[torch.Tensor], None] | None, force_sum_reduction_for_comms: bool = False, + label_suffix: str = "", ) -> tuple[ torch.Tensor, torch.Event, @@ -617,12 +630,15 @@ def foreach_reduce( ) _div_if_needed(reduce_scatter_input, predivide_factor) if world_size > 1: - reduce_scatter_comm( - output_tensor=reduce_output, - input_tensor=reduce_scatter_input, - group=reduce_scatter_group, - op=reduce_scatter_op, - ) + with dist.record_comm( + _label_with_suffix("FSDP::reduce_scatter", label_suffix) + ): + reduce_scatter_comm( + output_tensor=reduce_output, + input_tensor=reduce_scatter_input, + group=reduce_scatter_group, + op=reduce_scatter_op, + ) else: # For single GPU, just copy the input to output (no actual reduce-scatter needed), and # account for a possible gradient_divide_factor. @@ -655,11 +671,14 @@ def foreach_reduce( else: all_reduce_stream.wait_stream(current_stream) with device_handle.stream(all_reduce_stream): - dist.all_reduce( - reduce_output, - group=all_reduce_group, - op=all_reduce_op, - ) + with dist.record_comm( + _label_with_suffix("FSDP::all_reduce", label_suffix) + ): + dist.all_reduce( + reduce_output, + group=all_reduce_group, + op=all_reduce_op, + ) all_reduce_input = reduce_output all_reduce_event = all_reduce_stream.record_event() # -- END: ops in reduce_scatter stream @@ -811,7 +830,7 @@ def _get_gradient_divide_factors( if reduce_scatter_group is not None and factor == reduce_scatter_group.size(): reduce_scatter_op = ReduceOp.AVG else: - reduce_scatter_op = torch.distributed._make_nccl_premul_sum(1 / factor) + reduce_scatter_op = ReduceOp.PREMUL_SUM(1 / factor) return None, None, reduce_scatter_op, ReduceOp.SUM if factor is None: diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_param.py b/torch/distributed/fsdp/_fully_shard/_fsdp_param.py index 8b1227ee6d92e..f09e7b8dcd00a 100644 --- a/torch/distributed/fsdp/_fully_shard/_fsdp_param.py +++ b/torch/distributed/fsdp/_fully_shard/_fsdp_param.py @@ -238,6 +238,7 @@ def _init_sharded_param( else: self.mesh_info = mesh_info # pyrefly: ignore[bad-assignment] fsdp_placement = None + self._shard_mesh = self._init_shard_mesh() if param.device != device and param.device.type != "meta": raise AssertionError( f"Expects the parameter to already be moved to device {device} but got {param.device}" @@ -317,8 +318,10 @@ def _init_sharded_param( raise AssertionError( f"Expected contiguous tensor with {self.fsdp_placement=}" ) - self.sharded_param = nn.Parameter(self.to_sharded_dtensor(sharded_param)) - self.sharded_param.requires_grad_(param.requires_grad) + self.sharded_param = nn.Parameter( + self.to_sharded_dtensor(sharded_param), + requires_grad=param.requires_grad, + ) # Let `param_data` be freed normally when its ref count reaches 0 when # the `fully_shard` call returns to allow provided parameters to alias self._setattr_on_modules(self.sharded_param) @@ -529,8 +532,10 @@ def init_dtype_attrs(self, mp_policy: MixedPrecisionPolicy): # then we do not need extra casting if reduce_dtype == param_dtype: reduce_dtype = None - # Clamp `param_dtype` to `None` if no casting is required - if param_dtype == self.orig_dtype: + # Clamp `param_dtype` to `None` if no casting is required or if the + # parameter is non-floating-point (mixed precision is only meaningful + # for floating-point parameters) + if param_dtype == self.orig_dtype or not self.orig_dtype.is_floating_point: param_dtype = None self.param_dtype = param_dtype self.reduce_dtype = reduce_dtype @@ -659,7 +664,8 @@ def to_sharded_post_forward(self) -> None: storage_offset=0, ) self._sharded_post_forward_param = nn.Parameter( - self.to_sharded_post_forward_dtensor(sharded_post_forward_tensor) + self.to_sharded_post_forward_dtensor(sharded_post_forward_tensor), + requires_grad=self.sharded_param.requires_grad, ) self._setattr_on_modules(self._sharded_post_forward_param) self.free_unsharded_param() @@ -886,8 +892,7 @@ def _get_grad_inner_tensor(self, grad: torch.Tensor) -> torch.Tensor: def _sharded_local_tensor(self) -> torch.Tensor: return cast(DTensor, self.sharded_param)._local_tensor - @property - def shard_mesh(self): + def _init_shard_mesh(self) -> DeviceMesh: mesh = self.mesh_info.mesh if mesh.ndim == 1: return mesh @@ -895,6 +900,10 @@ def shard_mesh(self): raise AssertionError("Expected mesh_dim_names to not be None") return mesh[mesh.mesh_dim_names[-1]] + @property + def shard_mesh(self): + return self._shard_mesh + @property def shard_mesh_from_root(self): return self.shard_mesh diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py b/torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py index 0b71c63720e05..cd531a1a7f2b5 100644 --- a/torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py +++ b/torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py @@ -95,7 +95,7 @@ def lazy_init(self, device: torch.device): # tensors produced in one stream and used in another and accompanying # CUDA events for synchronization self.all_gather_state: AllGatherState | None = None - self.reduce_scatter_state: ReduceScatterState | None = None + self.reduce_scatter_states: list[ReduceScatterState] = [] # Post-forward order for explicit backward prefetching self.post_forward_order: list[FSDPParamGroup] = [] # will cause ref cycles @@ -189,6 +189,8 @@ def __init__( # - Communication and communication/computation overlap self.comm_ctx = FSDPCommContext() + self._param_group_index: int = 0 + self._num_param_groups: int = 1 # Group's indices in the shared post-forward order self._post_forward_indices: list[int] = [] # Whether to reduce gradients at all (whether for FSDP or HSDP) @@ -374,6 +376,7 @@ def unshard(self, async_op: bool = False): *self.comm_ctx.get_all_gather_streams(async_op, self._training_state), self.device, self._all_gather_comm, + self._label_suffix, ) def wait_for_unshard(self): @@ -546,17 +549,20 @@ def post_backward(self, *unused: Any): fsdp_param.unsharded_param.grad = None if self.reshard_after_backward: self.reshard() + # Wait on prior module's RS states (assumes backward fires groups + # N-1 first; if not, overlap degrades but correctness is preserved). + if ( + self._param_group_index == self._num_param_groups - 1 + and self.comm_ctx.reduce_scatter_states + ): + with record_function(f"FSDP::post_backward_rs_wait ({self._module_fqn})"): + for rs_state in self.comm_ctx.reduce_scatter_states: + if rs_state.event is not None: + self.device_handle.current_stream().wait_event(rs_state.event) + self.comm_ctx.reduce_scatter_states.clear() if len(fsdp_params_with_grad) == 0: return with record_function(self._with_fqn("FSDP::post_backward_reduce")): - if ( - self.comm_ctx.reduce_scatter_state is not None - and self.comm_ctx.reduce_scatter_state.event is not None - ): - self.device_handle.current_stream().wait_event( - self.comm_ctx.reduce_scatter_state.event - ) - self.comm_ctx.reduce_scatter_state = None all_reduce_pg = ( self._all_reduce_process_group if isinstance(self.mesh_info, DDPMeshInfo) @@ -607,9 +613,10 @@ def post_backward(self, *unused: Any): self._partial_reduce_output, self._all_reduce_hook, self.force_sum_reduction_for_comms, + self._label_suffix, ) - self.comm_ctx.reduce_scatter_state = ReduceScatterState( - reduce_scatter_input, reduce_scatter_event + self.comm_ctx.reduce_scatter_states.append( + ReduceScatterState(reduce_scatter_input, reduce_scatter_event) ) if all_reduce_input is not None: if self.device.type != "cpu": @@ -655,14 +662,36 @@ def _backward_prefetch(self) -> None: # Can be cleared if running multiple `backward`s return curr_index = self._post_forward_indices.pop() - if (target_index := curr_index - 1) < 0: - return - # Prefetch naively using the reverse post-forward order, which may - # have mistargeted prefetches if not all modules used in forward - # are used in this backward - # pyrefly: ignore [unbound-name] - target_fsdp_param_group = self.comm_ctx.post_forward_order[target_index] - self._prefetch_unshard(target_fsdp_param_group, "backward") + if self._num_param_groups > 1: + # Backward fires groups in reverse forward order: + # N-1, N-2, ..., 1, 0. Index 1 is always the + # penultimate group regardless of N. Prefetching here + # lets the next module's AG overlap with group 0's RS + # without holding unsharded params too long (as would + # happen if we prefetched from N-1). + if self._param_group_index != 1: + return + # E.g. fully_shard(block, shard_placement_fn=...) creates two + # param groups per block (dense + moe), giving + # post_forward_order = [block0, block0.moe, block1, block1.moe]. + # block1.moe walks back past block1 to prefetch block0.moe then block0. + curr_modules = self.modules + target_modules: tuple[nn.Module, ...] | None = None + for step in range(1, curr_index + 1): + target = self.comm_ctx.post_forward_order[curr_index - step] + if target.modules is curr_modules: + continue + if target_modules is None: + target_modules = target.modules + elif target.modules is not target_modules: + break + # Prefetch all groups of the target module in + # reverse forward order (highest index first), + # matching the explicit path in _pre_backward. + self._prefetch_unshard(target, "backward") + elif curr_index > 0: + target = self.comm_ctx.post_forward_order[curr_index - 1] + self._prefetch_unshard(target, "backward") @staticmethod def _prefetch_unshard( @@ -810,9 +839,17 @@ def _all_reduce_process_group(self) -> dist.ProcessGroup: ) return self.mesh_info.replicate_process_group + @property + def _label_suffix(self) -> str: + suffix = f"({self._module_fqn})" if self._module_fqn else "" + if self._num_param_groups > 1 and isinstance(self.mesh_info, FSDPMeshInfo): + suffix = f"{suffix} [pg={self.mesh_info.shard_mesh_size}]".lstrip() + return suffix + def _with_fqn(self, label: str) -> str: - if self._module_fqn: - return f"{label} ({self._module_fqn})" + suffix = self._label_suffix + if suffix: + return f"{label} {suffix}" return label def __repr__(self): diff --git a/torch/distributed/fsdp/_fully_shard/_fsdp_state.py b/torch/distributed/fsdp/_fully_shard/_fsdp_state.py index 90899abb5f3c5..a1bbcbf29c79f 100644 --- a/torch/distributed/fsdp/_fully_shard/_fsdp_state.py +++ b/torch/distributed/fsdp/_fully_shard/_fsdp_state.py @@ -227,8 +227,11 @@ def _init_shared_state(self) -> None: for state in self._state_ctx.all_states: state._state_ctx = self._state_ctx state._comm_ctx = self._comm_ctx - for fsdp_param_group in state._fsdp_param_groups: + num_groups = len(state._fsdp_param_groups) + for i, fsdp_param_group in enumerate(state._fsdp_param_groups): fsdp_param_group.comm_ctx = self._comm_ctx + fsdp_param_group._param_group_index = i + fsdp_param_group._num_param_groups = num_groups def _init_fqns(self) -> None: """Sets module and parameter FQN attributes for debugging.""" @@ -293,6 +296,8 @@ def _pre_forward( for fsdp_param_group in self._fsdp_param_groups: args, kwargs = fsdp_param_group.pre_forward(module, args, kwargs) for fsdp_state in self._states_to_forward_prefetch: + # Forward order (not reversed) to match forward execution order; + # contrast with reversed() in _pre_backward for backward order. for target_param_group in fsdp_state._fsdp_param_groups: FSDPParamGroup._prefetch_unshard(target_param_group, "forward") return args, kwargs @@ -333,7 +338,9 @@ def _pre_backward(self, grad: torch.Tensor) -> torch.Tensor: for fsdp_param_group in self._fsdp_param_groups: fsdp_param_group.pre_backward(default_prefetch) for fsdp_state in self._states_to_backward_prefetch: - for target_param_group in fsdp_state._fsdp_param_groups: + # Reverse so higher-indexed groups are prefetched first, + # matching backward execution order (reverse of forward). + for target_param_group in reversed(fsdp_state._fsdp_param_groups): FSDPParamGroup._prefetch_unshard(target_param_group, "backward") return grad @@ -342,7 +349,11 @@ def _root_post_backward_final_callback(self) -> None: logger.debug("FSDP::root_post_backward") with torch.profiler.record_function("FSDP::root_post_backward_callback"): for state in self._state_ctx.all_states: - for fsdp_param_group in state._fsdp_param_groups: + # Reverse so that the last param group (which gates the + # reduce-scatter wait/clear) fires first, matching the + # autograd backward order and preserving RS overlap for + # per-param-mesh modules whose inputs lack gradients. + for fsdp_param_group in reversed(state._fsdp_param_groups): if fsdp_param_group._training_state != TrainingState.POST_BACKWARD: # Run post-backward in case forward inputs did not require # gradient so the autograd backward did not run @@ -353,11 +364,12 @@ def _root_post_backward_final_callback(self) -> None: state._finalize_backward() if self._state_ctx.is_last_backward: self._comm_ctx.post_forward_order.clear() - if self._comm_ctx.reduce_scatter_state is not None: - self._device_handle.current_stream().wait_event( - self._comm_ctx.reduce_scatter_state.event - ) - self._comm_ctx.reduce_scatter_state = None + # Catch the last module's RS states that no subsequent + # module's group N-1 wait will clear. + for rs_state in self._comm_ctx.reduce_scatter_states: + if rs_state.event is not None: + self._device_handle.current_stream().wait_event(rs_state.event) + self._comm_ctx.reduce_scatter_states.clear() self._state_ctx.post_backward_final_callback_queued = False def _finalize_backward(self) -> None: diff --git a/torch/distributed/fsdp/_init_utils.py b/torch/distributed/fsdp/_init_utils.py index 633e743e7d5c8..349ccfd767207 100644 --- a/torch/distributed/fsdp/_init_utils.py +++ b/torch/distributed/fsdp/_init_utils.py @@ -617,13 +617,26 @@ def _init_param_handle_from_module( managed_params = list(_get_orig_params(fully_sharded_module, state._ignored_params)) _verify_managed_params(fully_sharded_module, managed_params) if sync_module_states: - _sync_module_params_and_buffers( - fully_sharded_module, managed_params, state.process_group - ) if state.sharding_strategy in HYBRID_SHARDING_STRATEGIES: + # Broadcast inter-node first, then intra-node. The inter-node + # broadcast propagates rank 0's states to each node's local + # rank 0, so the subsequent intra-node broadcast has the + # correct source values on every node. Reversing this order + # causes local rank 0 on non-source nodes to broadcast + # uninitialized states (e.g. from meta-device materialization). _sync_module_params_and_buffers( fully_sharded_module, managed_params, state._inter_node_pg ) + # _sync_module_params_and_buffers marks each buffer with + # FSDP_SYNCED=True to avoid redundant syncs in nested + # wrapping. Reset the flag here so the intra-node broadcast + # below also includes buffers. + for buffer in fully_sharded_module.buffers(): + if hasattr(buffer, FSDP_SYNCED): + setattr(buffer, FSDP_SYNCED, False) + _sync_module_params_and_buffers( + fully_sharded_module, managed_params, state.process_group + ) _init_param_handle_from_params(state, managed_params, fully_sharded_module) return state diff --git a/torch/distributed/fsdp/_runtime_utils.py b/torch/distributed/fsdp/_runtime_utils.py index 31c23b06705ac..4a5bd5f47211f 100644 --- a/torch/distributed/fsdp/_runtime_utils.py +++ b/torch/distributed/fsdp/_runtime_utils.py @@ -290,6 +290,7 @@ def _unshard( """ if not handle: return + handle._compute_stream = state._device_handle.current_stream() with state._device_handle.stream(pre_unshard_stream): ran_pre_unshard = handle.pre_unshard() if ran_pre_unshard: diff --git a/torch/distributed/launcher/api.py b/torch/distributed/launcher/api.py index 71ce57f62bce3..5e8140477b8a5 100644 --- a/torch/distributed/launcher/api.py +++ b/torch/distributed/launcher/api.py @@ -15,10 +15,17 @@ import torch import torch.distributed.elastic.rendezvous.registry as rdzv_registry -from torch._utils_internal import get_default_numa_options +from torch._utils_internal import get_default_numa_options, justknobs_check from torch.distributed.elastic import events, metrics from torch.distributed.elastic.agent.server.api import WorkerSpec -from torch.distributed.elastic.agent.server.local_elastic_agent import LocalElasticAgent +from torch.distributed.elastic.agent.server.health_check_server import ( + create_healthcheck_server, +) +from torch.distributed.elastic.agent.server.local_elastic_agent import ( + _AliveCallbackProxy, + LocalElasticAgent, + TORCHELASTIC_HEALTH_CHECK_PORT, +) from torch.distributed.elastic.multiprocessing import ( DefaultLogsSpecs, LogsSpecs, @@ -290,6 +297,34 @@ def launch_agent( # Set the signals to handle in the environment variable os.environ["TORCHELASTIC_SIGNALS_TO_HANDLE"] = config.signals_to_handle + # Start health check server before rendezvous so TW sees a healthy + # thrift port during the potentially long MAST rendezvous store barrier + # (10-22+ min for large jobs). The _AliveCallbackProxy returns + # time.time() until wired to the agent after construction. + health_check_server = None + alive_callback_proxy = None + healthcheck_port = os.getenv(TORCHELASTIC_HEALTH_CHECK_PORT) + if healthcheck_port is not None and justknobs_check( + "ai_infra/pytorch_distributed:torchelastic_enable_healthcheck_before_rendezvous", + default=False, + ): + try: + alive_callback_proxy = _AliveCallbackProxy() + health_check_server = create_healthcheck_server( + alive_callback=alive_callback_proxy, + port=int(healthcheck_port), + timeout=60, + ) + health_check_server.start() + logger.info( + "Started early health check server on port %s before rendezvous", + healthcheck_port, + ) + except Exception: + logger.warning("Failed to start early health check server", exc_info=True) + health_check_server = None + alive_callback_proxy = None + spec = WorkerSpec( role=config.role, local_world_size=config.nproc_per_node, @@ -314,8 +349,12 @@ def launch_agent( start_method=config.start_method, log_line_prefix_template=config.log_line_prefix_template, shutdown_timeout=config.shutdown_timeout, # type: ignore[arg-type] + health_check_server=health_check_server, ) + if alive_callback_proxy is not None: + alive_callback_proxy.set_delegate(agent._get_alive_time) + shutdown_rdzv = True try: metrics.initialize_metrics(metrics.MetricsConfig(config.metrics_cfg)) diff --git a/torch/distributed/pipelining/_backward.py b/torch/distributed/pipelining/_backward.py index c4e6256389aa3..7e8f418a02d4a 100644 --- a/torch/distributed/pipelining/_backward.py +++ b/torch/distributed/pipelining/_backward.py @@ -2,7 +2,7 @@ # Copyright (c) Meta Platforms, Inc. and affiliates import collections import logging -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from typing import Any import torch @@ -140,6 +140,40 @@ def get_param_groups( return unique_param_groups +def _autograd_grad_for_inputs( + outputs: Sequence[torch.Tensor], + inputs: Sequence[torch.Tensor], + grad_outputs: Sequence[torch.Tensor | None] | None = None, + retain_graph: bool = False, + allow_unused: bool = False, +) -> tuple[torch.Tensor | None, ...]: + """Compute input gradients, returning ``None`` for non-grad inputs.""" + # Some inputs may not be used or may not require gradients, so we filter them out + # before calling autograd.grad and place None for those positions in the result. + grad_indices: list[int] = [] + inputs_requiring_grad: list[torch.Tensor] = [] + for i, inp in enumerate(inputs): + if isinstance(inp, torch.Tensor) and inp.requires_grad: + grad_indices.append(i) + inputs_requiring_grad.append(inp) + + if not inputs_requiring_grad: + return tuple(None for _ in inputs) + + grads = torch.autograd.grad( + outputs=outputs, + inputs=inputs_requiring_grad, + grad_outputs=grad_outputs, + retain_graph=retain_graph, + allow_unused=allow_unused, + ) + + result: list[torch.Tensor | None] = [None] * len(inputs) + for idx, g in zip(grad_indices, grads, strict=True): + result[idx] = g + return tuple(result) + + def stage_backward_input( stage_outputs_or_loss: list[torch.Tensor], output_grads: list[torch.Tensor] | None, @@ -196,20 +230,20 @@ def hook(grad_inputs): torch.ones_like(stage_output) for stage_output in stage_outputs_or_loss ] - # Some inputs may not be used or may not require gradients, so we filter them out - input_values = [inp for inp in input_values if inp.requires_grad] - dinputs = torch.autograd.grad( + dinputs = _autograd_grad_for_inputs( stage_outputs_or_loss, - inputs=input_values, - grad_outputs=output_grads, + input_values, + output_grads, retain_graph=True, ) - # Update the gradients for inputs + + # Accumulate into .grad for inp, dinput in zip(input_values, dinputs): - if inp.grad is None: - inp.grad = dinput - else: - inp.grad += dinput + if isinstance(inp, torch.Tensor) and dinput is not None: + if inp.grad is None: + inp.grad = dinput + else: + inp.grad += dinput # stage_outputs_or_loss are not used in backwards after this point, so we can safely remove it from the autograd graph # this allows autograd to clear up the graph dedicated for this tensor and free up significant memory diff --git a/torch/distributed/pipelining/_utils.py b/torch/distributed/pipelining/_utils.py index 79b74be406814..3da67258ecbd9 100644 --- a/torch/distributed/pipelining/_utils.py +++ b/torch/distributed/pipelining/_utils.py @@ -1,91 +1,506 @@ # mypy: allow-untyped-defs # Copyright (c) Meta Platforms, Inc. and affiliates +from __future__ import annotations + import logging -from dataclasses import dataclass +import warnings +from dataclasses import dataclass, field +from enum import Enum +from typing import cast, Literal, overload, Protocol, TYPE_CHECKING, TypeAlias import torch from torch import fx +from torch.distributed._mesh_layout import _MeshLayout +from torch.distributed.tensor import DTensor +from torch.utils._pytree import tree_flatten, tree_unflatten + + +if TYPE_CHECKING: + from torch.distributed.device_mesh import DeviceMesh + from torch.distributed.tensor.placement_types import Placement logger = logging.getLogger(__name__) -def flatten_args_detach(args): +class GetMeshCallback(Protocol): + """Callback to create/retrieve a DeviceMesh from its cache key components.""" + + def __call__( + self, + mesh_dim_names: tuple[str, ...], + mesh_layout: _MeshLayout | None, + ) -> DeviceMesh: ... + + +# Key for mesh cache: (mesh_dim_names, mesh_layout) +# mesh_layout is the _MeshLayout object containing shape and stride (not actual ranks). +# This uniquely identifies a mesh within the same "universe" where all stages share +# the same rank tensor. +MeshCacheKey: TypeAlias = tuple[tuple[str, ...], _MeshLayout | None] + + +class PipeliningMetadataError(RuntimeError): + """Raised on metadata mismatches during pipeline communication.""" + + +@dataclass(frozen=True, slots=True) +class _TensorMeta: + """Tensor metadata for recv buffer allocation and validation. + + For plain tensors, these are the tensor's actual attributes. + For DTensors, these are LOCAL shard attributes; global attributes + are stored in :class:`_DTensorMeta`. """ - Flatten the args into a list form and detach the tensors from computational graph. + + shape: torch.Size + stride: tuple[int, ...] + dtype: torch.dtype + requires_grad: bool + + @staticmethod + def from_tensor(tensor: torch.Tensor) -> _TensorMeta: + """Create metadata from a plain tensor. + + Args: + tensor: A plain ``torch.Tensor`` (not DTensor). + + Returns: + Metadata capturing shape, stride, dtype, and requires_grad. + + Raises: + TypeError: If ``tensor`` is a DTensor. + """ + if isinstance(tensor, DTensor): + raise PipeliningMetadataError( + "Expected plain tensor, got DTensor. Use _DTensorMeta.from_dtensor instead." + ) + return _TensorMeta( + shape=tensor.shape, + stride=tensor.stride(), + dtype=tensor.dtype, + requires_grad=tensor.requires_grad, + ) + + def to_tensor(self, device: torch.device | str) -> torch.Tensor: + """Reconstruct a tensor on ``device`` from this metadata. + + Args: + device: Target device for the tensor. + + Returns: + An empty strided tensor on ``device``. + """ + t = _make_tensor_from_meta(self, device) + t.requires_grad_(self.requires_grad) + return t + + def get_diff(self, other: _TensorMeta) -> list[str]: + """Return field-by-field differences with ``other``. + + Args: + other: Metadata to compare against. + + Returns: + List of human-readable difference strings (empty if equal). + """ + if self == other: + return [] + + diffs = [] + if self.shape != other.shape: + diffs.append(f"shape mismatch: {self.shape} vs {other.shape}") + if self.stride != other.stride: + diffs.append(f"stride mismatch: {self.stride} vs {other.stride}") + if self.dtype != other.dtype: + diffs.append(f"dtype mismatch: {self.dtype} vs {other.dtype}") + # requires_grad is intentionally excluded: it is a runtime concern + # determined by has_backward and grad context, not a metadata invariant. + return diffs + + +@dataclass(frozen=True, slots=True) +class _DTensorMeta(_TensorMeta): + """DTensor metadata extending :class:`_TensorMeta` with distribution info. + + Inherited fields (shape, stride, etc.) are LOCAL shard attributes. + Additional fields capture global shape and placement information + needed to reconstruct a :class:`DTensor` via ``DTensor.from_local()``. + + The :class:`DeviceMesh` is **not** stored (not serializable for P2P); + it is looked up from :class:`_MeshCache` using + ``(mesh_dim_names, mesh_layout)`` as the key. """ - flat_detached_args = [] - def extract_tensor_args(a): - nonlocal flat_detached_args - if isinstance(a, torch.Tensor): - val = a.detach().requires_grad_(a.requires_grad) - flat_detached_args.append(val) - return val + # Global DTensor properties (for reconstruction) + global_shape: torch.Size = field(default_factory=lambda: torch.Size([])) + global_stride: tuple[int, ...] = field(default=()) + + # DTensor distribution properties + placements: tuple[Placement, ...] = field( + default=() + ) # e.g., (Shard(0), Replicate()) + + # Mesh identification - used to look up the correct DeviceMesh from cache + mesh_dim_names: tuple[str, ...] = field(default=()) # e.g., ("tp",) or ("dp", "tp") + mesh_layout: _MeshLayout | None = field( + default=None + ) # _MeshLayout with shape/stride - uniquely identifies mesh within the same universe + + @staticmethod + def from_dtensor(dtensor: DTensor) -> _DTensorMeta: + """Create metadata from a DTensor. + + Args: + dtensor: The DTensor to extract metadata from. + + Returns: + Metadata capturing both local and global attributes. + """ + device_mesh = dtensor.device_mesh + + return _DTensorMeta( + # Local tensor attributes (for recv buffer allocation) + shape=dtensor._local_tensor.shape, + stride=dtensor._local_tensor.stride(), + dtype=dtensor.dtype, + requires_grad=dtensor.requires_grad, + # Global DTensor attributes (for reconstruction) + global_shape=dtensor.shape, + global_stride=dtensor.stride(), + # Distribution info + placements=dtensor._spec.placements, + mesh_dim_names=( + tuple(device_mesh.mesh_dim_names) if device_mesh.mesh_dim_names else () + ), + mesh_layout=device_mesh._layout, + ) + + @property + def mesh_cache_key(self) -> MeshCacheKey: + """Cache key ``(mesh_dim_names, mesh_layout)`` for mesh lookup.""" + return (self.mesh_dim_names, self.mesh_layout) + + def to_dtensor(self, device: torch.device | str, mesh: DeviceMesh) -> DTensor: + """Reconstruct a DTensor on ``device`` with placements. + + Args: + device: Target device for the local tensor. + mesh: The ``DeviceMesh`` to attach. + + Returns: + A DTensor on ``device``. + """ + local_tensor = _make_tensor_from_meta(self, device) + # Set requires_grad after from_local() so that the from_local + # operation itself is not recorded in the autograd graph. + return cast( + DTensor, + DTensor.from_local( + local_tensor, + device_mesh=mesh, + placements=self.placements, + shape=self.global_shape, + stride=self.global_stride, + run_check=False, + ).requires_grad_(self.requires_grad), + ) + + def get_diff(self, other: _TensorMeta) -> list[str]: + """Return field-by-field differences, including DTensor-specific fields. + + Args: + other: Metadata to compare against. + + Returns: + List of human-readable difference strings (empty if equal). + """ + if self == other: + return [] + + # Get base class differences (compares local shape/stride/dtype/requires_grad) + # NOTE: Use explicit class call instead of super() because + # @dataclass(slots=True) on both parent and child can break super(). + diffs = _TensorMeta.get_diff(self, other) + + # Add DTensor-specific comparisons if other is also _DTensorMeta + if isinstance(other, _DTensorMeta): + if self.global_shape != other.global_shape: + diffs.append( + f"global_shape mismatch: {self.global_shape} vs {other.global_shape}" + ) + if self.global_stride != other.global_stride: + diffs.append( + f"global_stride mismatch: {self.global_stride} vs {other.global_stride}" + ) + if self.placements != other.placements: + diffs.append( + f"placements mismatch: {self.placements} vs {other.placements}" + ) + if self.mesh_dim_names != other.mesh_dim_names: + diffs.append( + f"mesh_dim_names mismatch: {self.mesh_dim_names} vs {other.mesh_dim_names}" + ) + if self.mesh_layout != other.mesh_layout: + diffs.append( + f"mesh_layout mismatch: {self.mesh_layout} vs {other.mesh_layout}" + ) else: - flat_detached_args.append(a) - return a + diffs.append("type: _DTensorMeta vs _TensorMeta") - new_args = fx.node.map_aggregate( - args, - extract_tensor_args, - ) + return diffs - return new_args, flat_detached_args +# Type alias for union of tensor metadata types +TensorMeta: TypeAlias = _TensorMeta | _DTensorMeta + + +# Not frozen: fields are populated incrementally during forward and +# backward metadata inference or from user provided static metadata +@dataclass(slots=True) +class _StageMeta: + """Consolidated tensor metadata for a pipeline stage's forward and backward passes.""" + + inputs: tuple[TensorMeta, ...] | None = None + outputs: tuple[TensorMeta, ...] | None = None + input_grads: tuple[TensorMeta | None, ...] | None = None + output_grads: tuple[TensorMeta | None, ...] | None = None + + def has_any(self) -> bool: + """Check if any metadata field is populated.""" + return any( + v is not None + for v in [self.inputs, self.outputs, self.input_grads, self.output_grads] + ) -def flatten_args(args): + def has_dtensors(self) -> bool: + """Check if any input/output metadata is DTensor type.""" + for metas in [self.inputs, self.outputs]: + if metas and any(isinstance(m, _DTensorMeta) for m in metas if m): + return True + return False + + def is_complete_for_forward(self) -> bool: + """Check if forward metadata is fully populated.""" + return self.inputs is not None and self.outputs is not None + + +@dataclass(frozen=True, slots=True) +class _StageForwardMeta: + """Forward metadata transmitted from stage *i* to stage *i+1* during inference.""" + + forward_metas: tuple[TensorMeta, ...] # Stage i's outputs → Stage i+1's inputs + + +@dataclass(frozen=True, slots=True) +class _StageBackwardMeta: + """Backward metadata transmitted from stage *i* to stage *i-1* during inference. + + Gradient placements may differ from forward activations + (e.g., ``Replicate`` → ``Partial``). """ - Flatten the args into a list form. + + backward_metas: tuple[ + TensorMeta | None, ... + ] # Stage i's input_grads → Stage i-1's output_grads + + +def _make_tensor_from_meta( + meta: _TensorMeta, + device: torch.device | str, +) -> torch.Tensor: + """Create a tensor from metadata. + + Args: + meta: Metadata with shape, stride, and dtype. + device: Target device for the tensor. + + Returns: + Empty tensor preserving the exact memory layout. """ - flat_args = [] + return torch.empty_strided( + size=meta.shape, + stride=meta.stride, + dtype=meta.dtype, + device=device, + ) + - def extract_tensor_args(a): - nonlocal flat_args - flat_args.append(a) - return a +def _derive_grad_metas( + tensor_metas: tuple[TensorMeta, ...], +) -> tuple[_TensorMeta | None, ...]: + """Derive gradient metadata from tensor metadata. - fx.node.map_aggregate( - args, - extract_tensor_args, + Returns metadata with the same shape/stride/dtype but ``requires_grad=False``. + Entries where the source has ``requires_grad=False`` become ``None``. + """ + return tuple( + _TensorMeta(shape=m.shape, stride=m.stride, dtype=m.dtype, requires_grad=False) + if m.requires_grad + else None + for m in tensor_metas ) - return flat_args +class _MeshCache: + """Cache for :class:`DeviceMesh` objects keyed by ``(mesh_dim_names, mesh_layout)``. -class PipeliningShapeError(RuntimeError): - """Shape mismatch between configured and runtime values.""" + Assumes all pipeline stages share the same rank tensor (true for + TorchTitan-style frameworks where meshes derive from a common world). + """ + def __init__(self, get_mesh_cb: GetMeshCallback | None = None) -> None: + self._cache: dict[MeshCacheKey, DeviceMesh] = {} + self._get_mesh_cb = get_mesh_cb -def validate_tensor_metadata(desc, expected, given): - if not expected.shape == given.shape: - raise PipeliningShapeError( - f"{desc} has a shape mismatch: expected {expected.shape} actual {given.shape}" - ) - if not expected.dtype == given.dtype: - raise PipeliningShapeError( - f"{desc} has a dtype mismatch: expected {expected.dtype} actual {given.dtype}" - ) - if not expected.stride() == given.stride(): - raise PipeliningShapeError( - f"{desc} has a stride mismatch: expected {expected.stride()} actual {given.stride()}" - ) + def get_mesh(self, key: MeshCacheKey) -> DeviceMesh: + """Return a cached mesh, or create one via the callback. + Args: + key: Cache key ``(mesh_dim_names, mesh_layout)``. -def validate_tensors_metadata( - desc, - expected_tensors: list[torch.Tensor] | tuple[torch.Tensor, ...], - actual_tensors: list[torch.Tensor] | tuple[torch.Tensor, ...], -): - if len(expected_tensors) != len(actual_tensors): - raise PipeliningShapeError( - f"{desc}: Number of values ({len(actual_tensors)}) does not match expected number ({len(expected_tensors)})" - ) - for i in range(len(expected_tensors)): - validate_tensor_metadata( - f"{desc}: value {i}", expected_tensors[i], actual_tensors[i] - ) + Returns: + The ``DeviceMesh``. + + Raises: + PipeliningMetadataError: If not cached and no callback provided. + """ + if key in self._cache: + return self._cache[key] + + mesh_dim_names, mesh_layout = key + + if self._get_mesh_cb is None: + raise PipeliningMetadataError( + f"Mesh not found in cache for mesh_dim_names={mesh_dim_names}, " + f"mesh_layout={mesh_layout}, and no get_mesh callback provided. " + f"Provide a get_mesh callback or use DTensors in static mode." + ) + + mesh = self._get_mesh_cb(mesh_dim_names, mesh_layout) + if mesh is None: + raise PipeliningMetadataError( + f"Mesh lookup failed: callback returned None for " + f"mesh_dim_names={mesh_dim_names}, mesh_layout={mesh_layout}. " + f"Ensure all stages use meshes from the same universe." + ) + self._cache[key] = mesh + return mesh + + def put(self, key: MeshCacheKey, mesh: DeviceMesh) -> None: + """Add a mesh to the cache.""" + self._cache[key] = mesh + + def update_from_tensors(self, tensors: tuple[torch.Tensor | None, ...]) -> None: + """Extract and cache meshes from any :class:`DTensor` instances in *tensors*.""" + for tensor in tensors: + if isinstance(tensor, DTensor): + mesh = tensor.device_mesh + dim_names = tuple(mesh.mesh_dim_names) if mesh.mesh_dim_names else () + mesh_layout = mesh._layout + key = (dim_names, mesh_layout) + if key not in self._cache: + self._cache[key] = mesh + + def __contains__(self, key: MeshCacheKey) -> bool: + return key in self._cache + + def __len__(self) -> int: + return len(self._cache) + + +# ============================================================================ +# Inference mode enum +# ============================================================================ + + +class InferenceMode(Enum): + """Pipeline-level metadata inference mode, determined collectively across all PP ranks. + + The mode is set by the schedule (not individual stages) because + ``has_backward`` is only known at schedule creation time and all + stages must agree to avoid P2P hangs. + + .. attribute:: STATIC + + All stages have sufficient metadata; runtime inference is skipped. + + .. attribute:: DYNAMIC + + At least one stage requires runtime metadata inference. + """ + + STATIC = "static" + DYNAMIC = "dynamic" + + @classmethod + def needs_dynamic(cls, meta: _StageMeta, stage_has_backward: bool) -> bool: + """Determine whether dynamic metadata inference is needed for a stage. + + Args: + meta: Stage metadata from user-provided args. + stage_has_backward: Whether a backward pass will be performed. + + Returns: + ``True`` if dynamic inference is needed. + """ + # Case 1: Forward metadata incomplete → needs DYNAMIC + if not meta.is_complete_for_forward(): + return True + + # Case 2: No DTensors → STATIC is fine (bwd metadata derivable from fwd metadata) + if not meta.has_dtensors(): + return False + + # Case 3: No backward needed → STATIC is fine (don't need grad metadata) + if not stage_has_backward: + return False + + # Case 4: DTensors with backward but missing ANY grad metadata → needs DYNAMIC + # Both input_grads AND output_grads are required for static mode with DTensors + if meta.input_grads is None or meta.output_grads is None: + return True + + # Case 5: DTensors with complete grads → STATIC is fine + return False + + +# ============================================================================ +# Utility functions +# ============================================================================ + + +def flatten_args(args, *, detach: bool = False): + """Flatten ``args`` into a list, optionally detaching tensors. + + Args: + args: Nested arguments to flatten. + detach: If ``True``, detach tensors while preserving ``requires_grad``. + + Returns: + ``(new_args, flat_detached_args)`` when ``detach=True``; + ``flat_args`` list otherwise. + """ + flat_args, treespec = tree_flatten(args) + + if detach: + flat_detached = [ + a.detach().requires_grad_(a.requires_grad) + if isinstance(a, torch.Tensor) + else a + for a in flat_args + ] + new_args = tree_unflatten(flat_detached, treespec) + return new_args, flat_detached + + return flat_args + + +# Backward compatibility alias +def flatten_args_detach(args): + """Flatten and detach. Deprecated: use ``flatten_args(args, detach=True)``.""" + return flatten_args(args, detach=True) def generate_stage_to_rank_mapping( @@ -148,7 +563,7 @@ def generate_rank_to_stage_mapping( return rank_to_stages -@dataclass +@dataclass(slots=True) class PipeInfo: """ Captures information for a pipeline (`Pipe` object). @@ -157,3 +572,367 @@ class PipeInfo: graph: fx.Graph num_stages: int has_loss_and_backward: bool + + +# ============================================================================ +# Metadata extraction helpers +# ============================================================================ + + +def extract_tensor_meta(tensor: torch.Tensor) -> TensorMeta: + """Extract metadata from a tensor. + + Handles both plain Tensor and DTensor correctly: DTensors are + dispatched to ``_DTensorMeta.from_dtensor`` which captures local + shard attributes plus global shape/placement info, while plain + tensors use ``_TensorMeta.from_tensor``. + + Args: + tensor: A plain tensor or DTensor. + + Returns: + ``_TensorMeta`` for plain tensors, ``_DTensorMeta`` for DTensors. + """ + if isinstance(tensor, DTensor): + return _DTensorMeta.from_dtensor(tensor) + else: + return _TensorMeta.from_tensor(tensor) + + +@overload +def extract_tensor_metas( + tensors: tuple[torch.Tensor, ...] | None, + *, + allow_none: Literal[False] = ..., +) -> tuple[TensorMeta, ...] | None: ... + + +@overload +def extract_tensor_metas( + tensors: tuple[torch.Tensor | None, ...] | None, + *, + allow_none: Literal[True], +) -> tuple[TensorMeta | None, ...] | None: ... + + +def extract_tensor_metas( + tensors: tuple[torch.Tensor | None, ...] | tuple[torch.Tensor, ...] | None, + *, + allow_none: bool = False, +) -> tuple[TensorMeta | None, ...] | None: + """Extract metadata from a tuple of tensors. + + Args: + tensors: Tuple of tensors (may include ``None`` when ``allow_none=True``). + allow_none: If ``True``, preserve ``None`` elements (for gradients). + + Returns: + Tuple of ``TensorMeta``, or ``None`` if ``tensors`` is ``None``. + + Raises: + PipeliningMetadataError: If ``None`` found and ``allow_none=False``. + """ + if tensors is None: + return None + + metas_with_none: list[TensorMeta | None] = [] + has_none = False + for t in tensors: + if isinstance(t, torch.Tensor): + metas_with_none.append(extract_tensor_meta(t)) + else: + has_none = True + metas_with_none.append(None) + if not allow_none and has_none: + raise PipeliningMetadataError( + "None values are not allowed in tensor metadata tuples. " + "Use allow_none=True for optional values." + ) + return tuple(metas_with_none) + + +def to_local_if_dtensor(tensor: torch.Tensor, detach: bool = False) -> torch.Tensor: + """Convert a DTensor to its local shard, or return a plain tensor as-is. + + When ``detach=True``, the tensor is detached before conversion — + this applies to both DTensors and plain tensors. + + Args: + tensor: A tensor that may be a DTensor. + detach: If ``True``, detach before ``to_local()`` to avoid + redistribution during backward. + + Returns: + The local tensor component. + """ + maybe_detached_tensor = tensor.detach() if detach else tensor + if isinstance(maybe_detached_tensor, DTensor): + return maybe_detached_tensor.to_local() + return maybe_detached_tensor + + +@overload +def validate_and_normalize_to_tuple( + args: torch.Tensor | tuple[torch.Tensor, ...] | list[torch.Tensor] | None, + allow_none: Literal[False] = ..., +) -> tuple[torch.Tensor, ...] | None: ... + + +@overload +def validate_and_normalize_to_tuple( + args: torch.Tensor + | tuple[torch.Tensor | None, ...] + | list[torch.Tensor | None] + | None, + allow_none: Literal[True] = ..., +) -> tuple[torch.Tensor | None, ...] | None: ... + + +def validate_and_normalize_to_tuple( + args: torch.Tensor + | tuple[torch.Tensor, ...] + | tuple[torch.Tensor | None, ...] + | list[torch.Tensor] + | list[torch.Tensor | None] + | None, + allow_none: bool = False, +) -> tuple[torch.Tensor | None, ...] | tuple[torch.Tensor, ...] | None: + """Normalize ``args`` to a tuple and validate that all elements are tensors. + + Args: + args: A single tensor, tuple/list of tensors, or ``None``. + allow_none: If ``True``, permit ``None`` elements (for gradients). + + Returns: + Tuple of tensors, or ``None`` if ``args`` is ``None``. + + Raises: + PipeliningMetadataError: On non-tensor values + (or ``None`` when ``allow_none=False``). + """ + if args is None: + return None + elif isinstance(args, torch.Tensor): + return (args,) + elif isinstance(args, (tuple, list)): + for i, arg in enumerate(args): + if arg is None: + if not allow_none: + raise PipeliningMetadataError( + f"Stage arg[{i}] is None. " + f"Stage args must be tensors. Use kwargs for optional values." + ) + continue + if not isinstance(arg, torch.Tensor): + raise PipeliningMetadataError( + f"Stage arg[{i}] has type {type(arg).__name__}. " + f"All stage args must be tensors. Use kwargs for non-tensor inputs." + ) + # Normalize list to tuple + return tuple(args) if isinstance(args, list) else args + else: + raise PipeliningMetadataError( + f"Stage args must be a tensor, tuple, or list of tensors, got {type(args).__name__}." + ) + + +# ============================================================================ +# Validation functions +# ============================================================================ + + +def validate_metadata( + desc: str, + expected: TensorMeta, + actual: torch.Tensor | TensorMeta, + *, + raise_on_mismatch: bool = False, + warn_on_mismatch: bool = False, +) -> list[str]: + """ + Compare expected metadata against actual tensor or metadata. + + This is the unified validation/comparison function that uses get_diff() from + metadata classes. Works with both plain tensors and DTensors. + + For plain tensors: compares shape/stride/dtype/requires_grad. + For DTensors: compares all properties including global shape and placements. + + Args: + desc: Description for error/warning messages. + expected: Expected tensor metadata (_TensorMeta or _DTensorMeta). + actual: Actual tensor or metadata to compare against. + raise_on_mismatch: If True, raise PipeliningMetadataError on mismatch. + warn_on_mismatch: If True, issue a warning on mismatch. + + Returns: + List of differences (empty if metadata matches). + + Raises: + PipeliningMetadataError: If raise_on_mismatch=True and differences exist. + """ + # Extract metadata if actual is a tensor + if isinstance(actual, torch.Tensor): + actual_meta = extract_tensor_meta(actual) + else: + actual_meta = actual + + # Type check: ensure both are same type for meaningful comparison + if type(expected) is not type(actual_meta): + type_diff = [ + f"type: expected {type(expected).__name__}, got {type(actual_meta).__name__}" + ] + if raise_on_mismatch: + raise PipeliningMetadataError(f"{desc}: {type_diff[0]}") + if warn_on_mismatch: + warnings.warn( + f"{desc}: Metadata type mismatch. {type_diff[0]}. " + f"Using dynamically inferred metadata instead.", + UserWarning, + stacklevel=2, + ) + return type_diff + + # Use get_diff() from the metadata class + diffs = expected.get_diff(actual_meta) + + if diffs: + if raise_on_mismatch: + raise PipeliningMetadataError(f"{desc}: {'; '.join(diffs)}") + if warn_on_mismatch: + warnings.warn( + f"{desc}: Metadata mismatch. {'; '.join(diffs)}. " + f"Using dynamically inferred metadata instead.", + UserWarning, + stacklevel=2, + ) + + return diffs + + +def validate_tensors_metadata( + desc: str, + expected: tuple[TensorMeta | None, ...], + actual: tuple[torch.Tensor | TensorMeta | None, ...], + *, + raise_on_mismatch: bool = True, + warn_on_mismatch: bool = False, +) -> list[str]: + """Validate metadata for a tuple of tensors element-wise. + + Args: + desc: Description prefix for error/warning messages. + expected: Tuple of expected metadata (may include ``None`` for grads). + actual: Tuple of actual tensors or metadata to compare against. + raise_on_mismatch: If ``True``, raise on the first mismatch. + warn_on_mismatch: If ``True``, issue warnings for mismatches. + + Returns: + Aggregated list of difference strings. + + Raises: + PipeliningMetadataError: If lengths differ or on mismatch. + """ + if len(expected) != len(actual): + msg = f"{desc}: expected {len(expected)} tensors, got {len(actual)}" + if raise_on_mismatch: + raise PipeliningMetadataError(msg) + if warn_on_mismatch: + warnings.warn(msg, UserWarning, stacklevel=2) + return [msg] + + all_diffs: list[str] = [] + for i, (exp, act) in enumerate(zip(expected, actual, strict=True)): + if exp is None and act is None: + continue + if exp is None or act is None: + msg = ( + f"{desc}[{i}]: expected {'None' if exp is None else 'metadata'}, " + f"got {'None' if act is None else 'metadata'}" + ) + if raise_on_mismatch: + raise PipeliningMetadataError(msg) + if warn_on_mismatch: + warnings.warn(msg, UserWarning, stacklevel=2) + all_diffs.append(msg) + continue + diffs = validate_metadata( + f"{desc}[{i}]", + exp, + act, + raise_on_mismatch=raise_on_mismatch, + warn_on_mismatch=warn_on_mismatch, + ) + all_diffs.extend(diffs) + return all_diffs + + +def validate_static_arg_grad_correspondence( + stage_index: int, + args: tuple[torch.Tensor, ...], + grads: tuple[torch.Tensor | None, ...], + is_input: bool, +) -> None: + """ + Validate the args↔grads contract for static mode. + + Enforces four rules for each (arg, grad) pair: + 1. len(args) must equal len(grads). + 2. If arg.requires_grad is False, grad must be None. + 3. If arg.requires_grad is True and grad is None, emit a warning + (this is legal at pipeline boundaries but may indicate a bug). + 4. If arg is a DTensor with requires_grad=True and grad is not None, + grad must also be a DTensor. + + Args: + stage_index: The stage index for error messages. + args: Tuple of forward tensors. + grads: Tuple of gradient tensors (can include None). + is_input: True for input_args/input_grads, False for output_args/output_grads. + + Raises: + PipeliningMetadataError: If any hard rule (1, 2, or 4) is violated. + """ + kind = "input" if is_input else "output" + args_name = f"{kind}_args" + grads_name = f"{kind}_grads" + + # Rule 1: lengths must match + if len(args) != len(grads): + raise PipeliningMetadataError( + f"Stage {stage_index}: {grads_name} length ({len(grads)}) does not match " + f"{args_name} length ({len(args)}). Each forward tensor must have a " + f"corresponding gradient entry (use None for tensors that don't require grad)." + ) + + for i, (arg, grad) in enumerate(zip(args, grads, strict=True)): + # Rule 2: no grad for a non-differentiable arg + if not arg.requires_grad and grad is not None: + raise PipeliningMetadataError( + f"Stage {stage_index}: {args_name}[{i}] has requires_grad=False, " + f"but {grads_name}[{i}] is not None ({type(grad).__name__}). " + f"Non-differentiable tensors must have None as their gradient entry." + ) + + # Rule 3: missing grad for a differentiable arg (warn, don't raise) + if arg.requires_grad and grad is None: + warnings.warn( + f"Stage {stage_index}: {args_name}[{i}] has requires_grad=True, " + f"but {grads_name}[{i}] is None. This is legal at pipeline boundaries " + f"but may indicate a missing gradient.", + UserWarning, + stacklevel=2, + ) + + # Rule 4: DTensor arg must have DTensor grad + if ( + isinstance(arg, DTensor) + and arg.requires_grad + and grad is not None + and not isinstance(grad, DTensor) + ): + raise PipeliningMetadataError( + f"Stage {stage_index}: {args_name}[{i}] is a DTensor with requires_grad=True, " + f"but {grads_name}[{i}] is {type(grad).__name__}, expected DTensor or None. " + f"DTensor gradients may have different placements than forward tensors." + ) diff --git a/torch/distributed/pipelining/microbatch.py b/torch/distributed/pipelining/microbatch.py index c1c353b191bb5..5b28c1942a59b 100644 --- a/torch/distributed/pipelining/microbatch.py +++ b/torch/distributed/pipelining/microbatch.py @@ -6,6 +6,8 @@ from typing import Any import torch +from torch.distributed.tensor import DTensor +from torch.distributed.tensor.experimental import local_map from torch.fx.node import map_aggregate from torch.nn.attention.flex_attention import BlockMask from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten @@ -202,26 +204,60 @@ def _split_tensor( raise AssertionError( f"Tensor size {tensor.size(spec.split_dim)} is smaller than num_chunks" ) - chunk_tensors = torch.tensor_split(tensor, num_chunks, spec.split_dim) - if not _debug_mask_minibatches: - return chunk_tensors - - expanded_chunks = [] - split_dim_idx = 0 - for chunk_tensor in chunk_tensors: - new_val = torch.zeros_like(tensor) - upper_idx = split_dim_idx + chunk_tensor.size(spec.split_dim) - - slice_indices = [slice(None, None, None)] * new_val.ndim - slice_indices[spec.split_dim] = slice(split_dim_idx, upper_idx) - new_val[slice_indices] = chunk_tensor + _is_dtensor = isinstance(tensor, DTensor) + + if _is_dtensor: + # Use local_map to split locally and preserve placements. + # Going through DTensor dispatch would convert Shard(split_dim) to + # Replicate() via an implicit all-gather, which is both wasteful and + # semantically wrong for PP microbatch splitting. + placements = tensor.placements + split_fn = local_map( + lambda t: torch.tensor_split(t, num_chunks, spec.split_dim), + out_placements=(placements,) * num_chunks, + in_placements=(placements,), + ) + chunk_tensors: Sequence[torch.Tensor] = split_fn(tensor) # type: ignore[assignment] + else: + chunk_tensors = torch.tensor_split(tensor, num_chunks, spec.split_dim) - expanded_chunks.append(new_val) + # tensor_split on a leaf tensor produces non-leaf views that won't + # accumulate .grad during torch.autograd.backward(). Call retain_grad() + # on those views so that stage_backward() can read .grad from them. + if tensor.requires_grad and tensor.is_leaf: + for chunk in chunk_tensors: + chunk.retain_grad() - split_dim_idx += chunk_tensor.size(spec.split_dim) + if not _debug_mask_minibatches: + return chunk_tensors - return expanded_chunks + def _expand_chunks( + orig: torch.Tensor, *chunks: torch.Tensor + ) -> tuple[torch.Tensor, ...]: + expanded = [] + idx = 0 + for chunk in chunks: + new_val = torch.zeros_like(orig) + upper = idx + chunk.size(spec.split_dim) + slices: list[slice] = [slice(None)] * new_val.ndim + slices[spec.split_dim] = slice(idx, upper) + new_val[slices] = chunk + expanded.append(new_val) + idx += chunk.size(spec.split_dim) + return tuple(expanded) + + if _is_dtensor: + placements = tensor.placements + n = len(chunk_tensors) + expand_fn = local_map( + _expand_chunks, + out_placements=(placements,) * n, + in_placements=(placements,) + (placements,) * n, + ) + return list(expand_fn(tensor, *chunk_tensors)) # type: ignore[arg-type] + else: + return list(_expand_chunks(tensor, *chunk_tensors)) def _shard_dict_of_args( @@ -538,7 +574,31 @@ def merge_chunks( else: values_to_cat = partial_values - args_flattened.append(torch.cat(values_to_cat, dim=arg.split_dim)) + # Validate DTensor consistency: either all values are DTensors + # or none are. A mix indicates a bug in the pipeline stage. + dtensor_flags = [isinstance(v, DTensor) for v in values_to_cat] + if any(dtensor_flags): + if not all(dtensor_flags): + raise AssertionError( + "merge_chunks: expected all values to be DTensors or " + "none to be DTensors, got a mix" + ) + # All DTensors must have matching placements. + placements = values_to_cat[0].placements + for i, v in enumerate(values_to_cat[1:], 1): + if v.placements != placements: + raise AssertionError( + f"merge_chunks: placement mismatch at chunk {i}: " + f"expected {placements}, got {v.placements}" + ) + cat_fn = local_map( + lambda *chunks: torch.cat(chunks, dim=arg.split_dim), + out_placements=(placements,), + in_placements=tuple(placements for _ in range(len(values_to_cat))), + ) + args_flattened.append(cat_fn(*values_to_cat)) + else: + args_flattened.append(torch.cat(values_to_cat, dim=arg.split_dim)) elif isinstance(arg, _CustomReducer): reduced_val = arg.init_value diff --git a/torch/distributed/pipelining/schedules.py b/torch/distributed/pipelining/schedules.py index 27789f74ba7a7..97ca83b019cc9 100644 --- a/torch/distributed/pipelining/schedules.py +++ b/torch/distributed/pipelining/schedules.py @@ -21,9 +21,13 @@ from torch.nn.modules.loss import _Loss from torch.profiler import record_function -from ._utils import generate_rank_to_stage_mapping, generate_stage_to_rank_mapping +from ._utils import ( + generate_rank_to_stage_mapping, + generate_stage_to_rank_mapping, + InferenceMode, +) from .microbatch import merge_chunks, split_args_kwargs_into_chunks, TensorChunkSpec -from .stage import _PipelineStageBase +from .stage import _PipelineStageBase, PipelineStage __all__ = [ @@ -320,6 +324,135 @@ def _update_losses(self, stages, losses): self._internal_losses.clear() + def _warmup_p2p( + self, + stages: list[_PipelineStageBase], + has_backward: bool, + p2p_done: bool, + ) -> None: + """Run the P2P warm-up protocol for the given stages. + + For ``PipelineStage`` instances this executes the forward/backward vote + protocol (which warms up 2-rank sub-communicators) and sets each + stage's ``_inference_mode``. For other stage types it falls back to + the legacy ``_get_init_p2p_neighbors_ops`` + ``_batch_p2p`` path. + + Args: + stages: The pipeline stages owned by this rank. + has_backward: Whether the schedule includes a backward pass. + p2p_done: ``True`` if P2P neighbours have already been initialised + (avoids redundant init on eval↔train mode switches). + """ + if all(isinstance(stage, PipelineStage) for stage in stages): + acc: torch.Tensor | None = None + for stage in cast(list[PipelineStage], stages): + acc = stage._warmup_forward_vote(has_backward, received_acc=acc) + result: torch.Tensor | None = acc + determined_mode: InferenceMode | None = None + for stage in reversed(cast(list[PipelineStage], stages)): + result = stage._warmup_backward_result(received_result=result) + if result is None: + raise RuntimeError("P2P warm-up voting failed") + determined_mode = ( + InferenceMode.STATIC + if result.item() == 1 + else InferenceMode.DYNAMIC + ) + stage._inference_mode = determined_mode + logger.debug( + "Rank determined inference_mode=%s for %d stage(s)", + determined_mode.value if determined_mode else "None", + len(stages), + ) + elif not p2p_done: + all_ops: list[dist.P2POp] = [] + for stage in stages: + all_ops.extend(stage._get_init_p2p_neighbors_ops()) + _wait_batch_p2p(_batch_p2p(all_ops)) + + # TODO: STATIC mode group communicator warm-up gap + # The vote protocol above warms up 2-rank sub-communicators + # (used by `_batch_p2p` homogeneous fast-path). In DYNAMIC mode, + # `_send_meta`/`_recv_meta` (called during `_prepare_forward_infra` → + # `_forward_metadata_inference`) also warm up the *group* communicator + # (used by `_batch_p2p` mixed-op path). In STATIC mode, metadata + # inference is skipped, so the group communicator is NOT warmed up — + # it will be lazily created on the first mixed `_batch_p2p` call + # (e.g., 1F1B steady-state with both sends and recvs). + # Fix: run `_get_init_p2p_neighbors_ops` + `_batch_p2p` after the + # vote, gated by `not p2p_done`. + + def _initialize_pp_stages( + self, + stages: list[_PipelineStageBase], + args: tuple[Any, ...] | Any, + kwargs: dict[str, Any] | None, + target: Any, + fwd_initialized: bool, + bwd_initialized: bool, + ) -> tuple[bool, bool]: + """Common stage initialization shared by Single and Multi schedules. + + Handles mode-change detection (eval↔train), P2P warm-up, RNG forking, + forward / backward metadata inference, and FSDP cleanup. + + Returns the updated ``(fwd_initialized, bwd_initialized)`` flags. + """ + # Detect eval↔train mode switch: if has_backward changed since last + # init, re-initialize both fwd (recv buffers need different + # requires_grad) and bwd. p2p_done avoids redundant P2P warm-up. + p2p_done = fwd_initialized + if fwd_initialized and (self._has_backward != bwd_initialized): + fwd_initialized = False + bwd_initialized = False + + needs_fwd = not fwd_initialized + needs_bwd = self._has_backward and not bwd_initialized + + if not needs_fwd and not needs_bwd: + return fwd_initialized, bwd_initialized + + if needs_fwd: + self._warmup_p2p(stages, self._has_backward, p2p_done) + + # Fork RNG so metadata inference doesn't perturb training RNG. + devices = list( + { + torch.device(stage.device) + for stage in stages + if torch.device(stage.device).type != "cpu" + } + ) + with torch.random.fork_rng(devices=devices): + if needs_fwd: + next_stage_args: Any = None + for stage in stages: + stage_args = args if stage.is_first else next_stage_args + next_stage_args = stage._prepare_forward_infra( + self._n_microbatches, + stage_args, + kwargs, + has_backward=self._has_backward, + ) + fwd_initialized = True + + if needs_bwd: + prev_stage_grad_meta: Any = None + for stage in reversed(stages): + prev_stage_grad_meta = stage._prepare_backward_infra( + self._n_microbatches, + loss_fn=self._loss_fn, + target=target, + received_grad_meta=prev_stage_grad_meta, + ) + bwd_initialized = True + + for stage in stages: + if isinstance(stage, PipelineStage): + stage._post_metadata_inference_cleanup() + + return fwd_initialized, bwd_initialized + @abstractmethod def _step_microbatches( self, @@ -560,21 +693,18 @@ def __init__( self._get_pipeline_order() ) - def _initialize_stage(self, args, kwargs): - if not self._stage_forward_initialized: - # Prepare the communication needed for the pipeline schedule execution - # This is needed because during execution we always perform a series of batch P2P ops - # The first call of the batched P2P needs to involve the global group - all_ops: list[dist.P2POp] = [] - all_ops.extend(self._stage._get_init_p2p_neighbors_ops()) - _wait_batch_p2p(_batch_p2p(all_ops)) - - self._stage._prepare_forward_infra(self._n_microbatches, args, kwargs) - self._stage_forward_initialized = True - - if self._has_backward and not self._stage_backward_initialized: - self._stage._prepare_backward_infra(self._n_microbatches) - self._stage_backward_initialized = True + def _initialize_stage(self, args, kwargs, target=None): + ( + self._stage_forward_initialized, + self._stage_backward_initialized, + ) = self._initialize_pp_stages( + [self._stage], + args, + kwargs, + target, + self._stage_forward_initialized, + self._stage_backward_initialized, + ) def step( self, @@ -670,7 +800,8 @@ def _step_microbatches( ) arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses) - self._initialize_stage(arg_mbs[0], kwarg_mbs[0]) + maybe_first_target = target_mbs[0] if target_mbs is not None else None + self._initialize_stage(arg_mbs[0], kwarg_mbs[0], maybe_first_target) # Delay send waits fwd_sends_to_wait: list[list[dist.Work]] = [] @@ -721,7 +852,8 @@ def _step_microbatches( return_outputs: whether to return the outputs from the last stage. """ arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses) - self._initialize_stage(arg_mbs[0], kwarg_mbs[0]) + maybe_first_target = target_mbs[0] if target_mbs is not None else None + self._initialize_stage(arg_mbs[0], kwarg_mbs[0], maybe_first_target) # Delay send waits fwd_sends_to_wait: list[list[dist.Work]] = [] @@ -865,7 +997,8 @@ def _step_microbatches( return_outputs: whether to return the outputs from the last stage. """ arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses) - self._initialize_stage(arg_mbs[0], kwarg_mbs[0]) + maybe_first_target = target_mbs[0] if target_mbs is not None else None + self._initialize_stage(arg_mbs[0], kwarg_mbs[0], maybe_first_target) # Last stage has 1 warmup, second-to-last 2 warmups, ... # first stage `num_stages` warmups @@ -1532,34 +1665,18 @@ def __init__( "Simply stop passing it, and everything should still work fine." ) - def _initialize_stages(self, args: tuple[Any, ...], kwargs): - if not self._stages_forward_initialized: - # Prepare the communication needed for the pipeline schedule execution - # This is needed because during execution we always perform a series of batch P2P ops - # The first call of the batched P2P needs to involve the global group - all_ops: list[dist.P2POp] = [] - for stage in self._stages: - all_ops.extend(stage._get_init_p2p_neighbors_ops()) - _wait_batch_p2p(_batch_p2p(all_ops)) - - # may be 'none' value (if this stage sends its output shapes to the next stage via P2P) - # or real value (if this stage and next stage are on the same device) - next_stage_args: tuple[Any, ...] = tuple() - for stage in self._stages: - if stage.is_first: - next_stage_args = stage._prepare_forward_infra( - self._n_microbatches, args, kwargs - ) - else: - next_stage_args = stage._prepare_forward_infra( - self._n_microbatches, next_stage_args, kwargs - ) - self._stages_forward_initialized = True - - if self._has_backward and not self._stages_backward_initialized: - for stage in self._stages: - stage._prepare_backward_infra(self._n_microbatches) - self._stages_backward_initialized = True + def _initialize_stages(self, args: tuple[Any, ...], kwargs, target=None): + ( + self._stages_forward_initialized, + self._stages_backward_initialized, + ) = self._initialize_pp_stages( + self._stages, + args, + kwargs, + target, + self._stages_forward_initialized, + self._stages_backward_initialized, + ) def _validate_and_set_stage_mapping( self, actions: dict[int, list[_Action | None]] @@ -1674,8 +1791,8 @@ def _step_microbatches( not support models with skip connections. """ arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses) - - self._initialize_stages(arg_mbs[0], kwarg_mbs[0]) + maybe_first_target = target_mbs[0] if target_mbs is not None else None + self._initialize_stages(arg_mbs[0], kwarg_mbs[0], maybe_first_target) # Based on the plan in Step 1 created in __init__: # 2. Perform communication based on the pipeline_order @@ -1839,7 +1956,7 @@ def _step_microbatches( at time_step %s when running action %s", self.rank, self.__class__.__name__, - str(e), + e, time_step, action, ) @@ -2060,7 +2177,8 @@ def _step_microbatches( not support models with skip connections. """ arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses) - self._initialize_stages(arg_mbs[0], kwarg_mbs[0]) + maybe_first_target = target_mbs[0] if target_mbs is not None else None + self._initialize_stages(arg_mbs[0], kwarg_mbs[0], maybe_first_target) # Based on the plan in Step 1 created in __init__: # 2. Perform communication based on the pipeline_order diff --git a/torch/distributed/pipelining/stage.py b/torch/distributed/pipelining/stage.py index cd56efbbf4420..0fbe780714869 100644 --- a/torch/distributed/pipelining/stage.py +++ b/torch/distributed/pipelining/stage.py @@ -2,6 +2,7 @@ # Copyright (c) Meta Platforms, Inc. and affiliates import logging import operator +import warnings from abc import ABC, abstractmethod from collections.abc import Callable from typing import Any, cast @@ -13,13 +14,39 @@ from torch._subclasses.fake_tensor import FakeTensor from torch.distributed._composable.replicate_with_fsdp import replicate, ReplicateModule from torch.distributed.fsdp import FSDPModule, fully_shard +from torch.distributed.pipelining._utils import ( + _derive_grad_metas, + _DTensorMeta, + _make_tensor_from_meta, + _MeshCache, + _StageBackwardMeta, + _StageForwardMeta, + _StageMeta, + _TensorMeta, + extract_tensor_meta, + extract_tensor_metas, + flatten_args, + GetMeshCallback, + InferenceMode, + PipeInfo, + PipeliningMetadataError, + TensorMeta, + to_local_if_dtensor, + validate_and_normalize_to_tuple, + validate_static_arg_grad_correspondence, + validate_tensors_metadata, +) +from torch.distributed.tensor import DTensor from torch.fx.node import Argument, map_aggregate from torch.nn.parallel import DistributedDataParallel -from torch.utils._pytree import tree_map_only -from ._backward import stage_backward, stage_backward_input, stage_backward_weight +from ._backward import ( + _autograd_grad_for_inputs, + stage_backward, + stage_backward_input, + stage_backward_weight, +) from ._debug import map_debug_info -from ._utils import flatten_args, PipeInfo, validate_tensors_metadata __all__ = [ @@ -63,61 +90,47 @@ def _normalize_model_output_as_tuple(output: Any) -> tuple[Any]: return output_tuple -class _RootArgPlaceholder: - """ - Placeholder for model-level inputs. - """ - - def __init__(self, tensor): - self.meta = tensor.to("meta") - - class _RecvInfo: - """ - Represents a stage input. + """Input tensor descriptor for a pipeline stage. + + Handles both received activations from a previous stage + (``is_root_arg=False``) and root-level model inputs provided + by the user (``is_root_arg=True``). """ def __init__( self, input_name: str, - source: int, - buffer: torch.Tensor, + source: int | None, + buffer: torch.Tensor | None, + tensor_meta: TensorMeta | None, + *, + is_root_arg: bool = False, ): # Name of this input self.input_name = input_name - # Stage index of the source of this input + # Stage index of the source of this input (None for root args) self.source = source - # Buffer to receive the input into. + # Buffer to receive the input into (None for root args) self.buffer = buffer + # Tensor metadata for validation and DTensor reconstruction + self.tensor_meta = tensor_meta + # Whether this is a root-level model input (no recv needed) + self.is_root_arg = is_root_arg def __repr__(self): - return f"_RecvInfo(input={self.input_name}, source={self.source}, shape={self.buffer.size()})" - - -# An input can be either a received activation or a model input -InputInfo = _RecvInfo | _RootArgPlaceholder - - -def _make_tensor_from_meta( - example: torch.Tensor | FakeTensor, - device: torch.device, -) -> torch.Tensor: - """ - Create a real tensor from a tensor. - """ - return torch.empty( - example.size(), - dtype=example.dtype, - layout=example.layout, - device=device, - ) + if self.is_root_arg: + return f"_RecvInfo(input={self.input_name}, root_arg=True)" + meta_type = type(self.tensor_meta).__name__ if self.tensor_meta else "None" + buffer_shape = self.buffer.size() if self.buffer is not None else "None" + return f"_RecvInfo(input={self.input_name}, source={self.source}, shape={buffer_shape}, meta={meta_type})" class _PipelineStageBase(ABC): - """ - Base class for pipeline stages. - Defines or implements common methods used by the `_PipelineStage` used by - the tracing frontend and `PipelineStage` used by manual frontend. + """Base class for pipeline stages. + + Defines common methods used by ``_PipelineStage`` (tracing frontend) + and ``PipelineStage`` (manual frontend). """ def __init__( @@ -131,20 +144,16 @@ def __init__( ): """ Args: - submodule (torch.nn.Module): The module to be executed in this stage. - stage_index (int): The index of this stage. - num_stages (int): The total number of stages in this pipeline. - device (torch.device): The device to run this stage on. - group (Optional[dist.ProcessGroup]): The process group to use for communication. - If `None`, the default process group will be used. - Default: `None`. - dw_builder (Optional[Callable[[], Callable[..., None]]): If provided, dw_builder is a builder function - that will build a new dw_runner function that will run parts of module backward that were intentionally - skipped during the module's actual backward pass. The builder must be invoked by stage after stage runs - model backwards, and stage should save the latest dw_runner to run during weight pas (W). - If not provided, a dw_runner will be generated automatically by traversing the autograd graph. - When used with schedules that only have F and B steps, the fresh dw_runner function will be called as - part of I (input backwards). When used with F,I,W schedules, the dw_runner function implements 'W'. + submodule: The module to be executed in this stage. + stage_index: The index of this stage. + num_stages: The total number of stages in this pipeline. + device: The device to run this stage on. + group: Process group for communication. Defaults to the + default process group if ``None``. + dw_builder: Builder function that produces a ``dw_runner`` + for deferred weight updates in F/I/W zero-bubble + schedules. If ``None``, a runner is generated + automatically via autograd graph traversal. """ super().__init__() if stage_index >= num_stages: @@ -175,7 +184,6 @@ def __init__( ) # Run time states - self._outputs_meta: tuple[torch.Tensor, ...] | None = None # map microbatch ID to list of forward tensor args self.fwd_cache: dict[int, tuple[Any, list[torch.Tensor]]] = {} # map microbatch ID to list of backward grad tensor args @@ -190,7 +198,7 @@ def __init__( self.log_prefix = f"[Stage {self.stage_index}]" # Forward infra - self.args_recv_info: dict[int, tuple[InputInfo, ...]] = {} + self.args_recv_info: dict[int, tuple[_RecvInfo, ...]] = {} self.act_send_info: dict[int, list] = {} # Backward infra will created lazily @@ -203,6 +211,17 @@ def __init__( i: i % self.group_size for i in range(self.num_stages) } + # DTensor support: mesh cache for looking up DeviceMesh by (dim_names, layout) + self._mesh_cache = _MeshCache() + + # Per-chunk runtime validation is expensive; only enable under + # TORCH_DISTRIBUTED_DEBUG=DETAIL for debugging shape/dtype mismatches. + self._runtime_validate = dist.get_debug_level() == dist.DebugLevel.DETAIL + + # DTensor support: consolidated stage metadata container + # Contains inputs, outputs, input_grads, output_grads metadata + self._stage_meta = _StageMeta() + @property def has_backward(self) -> bool: """ @@ -228,6 +247,21 @@ def is_last(self): """ return self.stage_index == self.num_stages - 1 + def _validate_stage_tensors( + self, + desc: str, + expected: tuple[TensorMeta | None, ...] | None, + actual: tuple[torch.Tensor | None, ...], + ) -> None: + """Validate actual tensors against expected metadata. + + Raises: + PipeliningMetadataError: If metadata is missing or mismatched. + """ + if expected is None: + raise PipeliningMetadataError(f"{desc}: no metadata available") + validate_tensors_metadata(desc, expected, actual) + def _check_chunk_id(self, chunk_id: int): if self.chunks is None: raise RuntimeError( @@ -238,27 +272,6 @@ def _check_chunk_id(self, chunk_id: int): f"Chunk id {chunk_id} is out of range [0, {self.chunks})" ) - def _configure_outputs_meta(self, outputs_meta: tuple[torch.Tensor, ...]): - """ - Track the output shapes/dtype of this stage since they determine the send operation(s) which must match - recv operations of the next stage. The next stage _will_ be freezing its recv buffers based on its initial - configuration, so it's important to also freeze/validate the output side to avoid any send/recv mismatches - which could show up as hangs, silent corruption, or other errors. - """ - if self._outputs_meta is not None: - raise AssertionError( - "Attempting to reconfigure output_meta, which is not supported" - ) - self._outputs_meta = tuple(outputs_meta) # type: ignore[assignment] - - def get_outputs_meta(self) -> tuple[torch.Tensor, ...]: - """Get the output metadata (meta tensors) reprensenting the outputs of this stage""" - if self._outputs_meta is None: - raise AssertionError( - "Attempted to get_outputs_meta() without configuring output meta" - ) - return self._outputs_meta - def _create_grad_send_info( self, args_recv_info: tuple, @@ -272,12 +285,13 @@ def map_recv_to_send(a): # Note: we send gradients back to previous stage as long as in # forward it is a received input, regardless of whether it requires # grad. It is up to the previous stage to discard this gradient. - if isinstance(a, _RecvInfo): - grad_send_info.append(a.source) - return a.source - else: + if a.is_root_arg: + # Root args don't have a source stage to send gradients to grad_send_info.append(None) return None + else: + grad_send_info.append(a.source) + return a.source map_aggregate(args_recv_info, map_recv_to_send) @@ -288,17 +302,30 @@ def map_recv_to_send(a): def _prepare_forward_infra( self, num_microbatches: int, - args: tuple[Any, ...], + args: tuple[Any, ...] | _StageForwardMeta | None, kwargs: dict[str, Any] | None = None, - ) -> tuple[Any, ...]: + has_backward: bool = False, + ) -> _StageForwardMeta | None: + raise NotImplementedError + + @abstractmethod + def _prepare_backward_infra( + self, + num_microbatches: int, + loss_fn: Callable[..., torch.Tensor] | None = None, + target: torch.Tensor | None = None, + received_grad_meta: _StageBackwardMeta | None = None, + ) -> _StageBackwardMeta | None: raise NotImplementedError - def _prepare_backward_infra(self, num_microbatches: int): + def _setup_backward_recv_info(self, num_microbatches: int): # TODO: this is needed for backward_maybe_with_nosync self.chunks = num_microbatches + # IMPORTANT: _create_grad_recv_info reads self._stage_meta.output_grads + # to attach DTensor metadata to _RecvInfo objects. The clear below MUST + # happen after all _create_grad_recv_info calls complete. for mb_index in range(num_microbatches): - # `grad_recv_info` is a mirror of `act_send_info` self.grad_recv_info[mb_index] = self._create_grad_recv_info( self.act_send_info ) @@ -310,9 +337,17 @@ def _create_grad_recv_info( ) -> tuple[_RecvInfo, ...]: raise NotImplementedError + def _resolve_peer_global_rank(self, stage_idx: int) -> int: + """Map a pipeline stage index to the corresponding global rank for P2P communication.""" + peer_rank = self.stage_index_to_group_rank[stage_idx] + return dist.get_global_rank( + self.group or dist.distributed_c10d._get_default_group(), + peer_rank, + ) + def _get_recv_ops( self, - recv_infos: tuple[InputInfo, ...], + recv_infos: tuple[_RecvInfo, ...], ) -> list[dist.P2POp]: """ Helper function shared by `get_fwd_recv_ops` and `get_bwd_recv_ops`. @@ -320,15 +355,16 @@ def _get_recv_ops( """ ops: list[dist.P2POp] = [] for info in recv_infos: - if not isinstance(info, _RecvInfo): + if info.is_root_arg: + # Root args don't need recv operations continue - - peer_rank = self.stage_index_to_group_rank[info.source] - peer_global_rank = ( - peer_rank - if self.group is None - else dist.get_global_rank(self.group, peer_rank) - ) + # Skip entries with None buffer (None gradients) + if info.buffer is None: + assert info.tensor_meta is None # noqa: S101 + continue + # At this point, source and buffer are guaranteed non-None + assert info.source is not None # noqa: S101 + peer_global_rank = self._resolve_peer_global_rank(info.source) ops.append( dist.P2POp(dist.irecv, info.buffer, peer_global_rank, self.group) ) @@ -337,46 +373,42 @@ def _get_recv_ops( """[Note: V-schedule special case] - V-Schedules have a special case where 2 stages with adjacent stage_id are on the same rank. + V-Schedules have a special case where 2 stages with adjacent stage_id + are on the same rank. - ex: 2 ranks, 4 stages forms a simple V: - rank0: stage 0 stage 3 - rank1: stage 1 stage 2 + Example: 2 ranks, 4 stages forms a simple V:: - stage 0,1 and 2,3 communicate activations using send/recv as usual, but stage 1,2 do not need to - use communication ops. Instead, they should pass tensor data directly via function call. + rank0: stage 0 stage 3 + rank1: stage 1 stage 2 - set_local_fwd_input and (get_local_bwd_output + set_local_bwd_input) facilitate this optimization, and - should be called at the appropriate time during the pipeline schedule (after forward or backward execution). + Stages 0/1 and 2/3 communicate via send/recv, but stages 1/2 pass + tensors directly via function call, avoiding communication ops. """ def set_local_fwd_input(self, prev_stage_outputs: Any, mb_index: int) -> None: + """Pass outputs from a same-rank stage as forward inputs (V-schedule). + + Detaches tensors and sets ``requires_grad`` so they serve as autograd + leaves. Handles DTensor activations transparently. """ - Moves 'prev_stage_outputs' from another stage on the same rank into place as inputs for this stage. Avoids - copying tensor data or using send/recv op. Detaches original tensor and sets requires_grad so the - tensor can serve as a leaf for autograd and gradients can be collected from it during backward. - """ - recv_infos: tuple[InputInfo, ...] = self.args_recv_info[mb_index] + recv_infos: tuple[_RecvInfo, ...] = self.args_recv_info[mb_index] # See [Note: pipeline model output type] prev_stage_outputs = _normalize_model_output_as_tuple(prev_stage_outputs) - for info, tensor in zip(recv_infos, prev_stage_outputs): + for info, tensor in zip(recv_infos, prev_stage_outputs, strict=True): if not isinstance(tensor, torch.Tensor): raise AssertionError( f"expected tensor values as outputs from prev stage, got {type(tensor)}" ) - if not isinstance(info, _RecvInfo): + if info.is_root_arg: raise AssertionError( - "set_local_Fwd_input should only be called on non-first stage, which should always have RecvInfo" + "set_local_fwd_input should only be called on non-first stage, which should always have non-root RecvInfo" ) - # We don't need to do a data copy here, since we can directly pass the activation tensor reference from - # one stage to the next. However, we do need to mark the activation as a leaf tensor since it will serve - # as the input tensor for a fresh autograd graph, not part of the previous stage's autograd graph. - # TODO: confirm, do we use this activation as the root of the backward call for the previous stage? does - # detach have any affect on that? - info.buffer = tensor.detach().requires_grad_(True) + # Pass the activation tensor directly (same rank for local execution). + # Detach to create a new autograd leaf for the fresh autograd graph. + info.buffer = to_local_if_dtensor(tensor, detach=True) def get_local_bwd_output(self, mb_index): """ @@ -398,6 +430,7 @@ def set_local_bwd_input( """ Moves 'grad input' tensors from the next stage to 'grad_output' on this stage, avoiding a copy or send/recv. Does not detach or set '_requires_grad'. + Handles DTensor gradients for V-schedule local passing. """ if not isinstance(next_stage_bwd_outputs, tuple): raise AssertionError(f"Expected tuple, got {type(next_stage_bwd_outputs)}") @@ -409,21 +442,27 @@ def set_local_bwd_input( if self.is_last: raise AssertionError("can't set bwd input if this stage is last") recv_infos = self.grad_recv_info[mb_index] - for info, tensor in zip(recv_infos, next_stage_bwd_outputs): + for info, tensor in zip(recv_infos, next_stage_bwd_outputs, strict=True): + if tensor is None: + continue if not isinstance(tensor, torch.Tensor): raise AssertionError( f"expected tensor values as outputs from prev stage, got {type(tensor)}" ) - if not isinstance(info, _RecvInfo): - raise AssertionError(f"Expected a recv info, got {type(info)}") - info.buffer = tensor + if info.is_root_arg: + raise AssertionError( + "set_local_bwd_input should only be called with non-root RecvInfo" + ) + + # Extract local tensor for the buffer (handles DTensor or plain tensor) + info.buffer = to_local_if_dtensor(tensor) def get_fwd_recv_ops(self, fwd_chunk_id: int) -> list[dist.P2POp]: """ Returns a list of ops that are needed to receive the input arguments for this stage. """ - recv_infos: tuple[InputInfo, ...] = self.args_recv_info[fwd_chunk_id] + recv_infos: tuple[_RecvInfo, ...] = self.args_recv_info[fwd_chunk_id] return self._get_recv_ops(recv_infos) @@ -441,6 +480,7 @@ def get_bwd_recv_ops(self, bwd_chunk_id: int) -> list[dist.P2POp]: def get_fwd_send_ops(self, fwd_chunk_id: int) -> list[dist.P2POp]: """ Get the activation send ops for current stage's forward. + Handles DTensor outputs by extracting local tensors. """ output_tuple, _ = self.fwd_cache[fwd_chunk_id] @@ -451,25 +491,25 @@ def get_fwd_send_ops(self, fwd_chunk_id: int) -> list[dist.P2POp]: for dst in dst_stages: if dst is None: continue + # Extract local tensor if DTensor + send_tensor = to_local_if_dtensor(out, detach=True) logger.debug( "%s Sending tensor to Stage %s: %s", self.log_prefix, dst, - out.size(), + send_tensor.size(), ) - peer_rank = self.stage_index_to_group_rank[dst] - peer_global_rank = ( - peer_rank - if self.group is None - else dist.get_global_rank(self.group, peer_rank) + peer_global_rank = self._resolve_peer_global_rank(dst) + ops.append( + dist.P2POp(dist.isend, send_tensor, peer_global_rank, self.group) ) - ops.append(dist.P2POp(dist.isend, out, peer_global_rank, self.group)) return ops def get_bwd_send_ops(self, bwd_chunk_id: int) -> list[dist.P2POp]: """ Get the gradient send ops for current stage's backward. + Handles DTensor gradients by extracting local tensors. """ if not self.has_backward or self.is_first: return [] @@ -485,24 +525,24 @@ def get_bwd_send_ops(self, bwd_chunk_id: int) -> list[dist.P2POp]: ops: list[dist.P2POp] = [] grads_input = self.bwd_cache.pop(bwd_chunk_id) - for grad, grad_recv_stage in zip(grads_input, self.grad_send_info): + + for grad, grad_recv_stage in zip(grads_input, self.grad_send_info, strict=True): if isinstance(grad, torch.Tensor) and grad_recv_stage is not None: + # Extract local tensor if DTensor + send_tensor = to_local_if_dtensor(grad) logger.debug( "%s Sending gradient to Stage %s: %s", self.log_prefix, grad_recv_stage, - grad.size(), + send_tensor.size(), ) - peer_rank = self.stage_index_to_group_rank[grad_recv_stage] - peer_global_rank = ( - peer_rank - if self.group is None - else dist.get_global_rank(self.group, peer_rank) + peer_global_rank = self._resolve_peer_global_rank(grad_recv_stage) + ops.append( + dist.P2POp(dist.isend, send_tensor, peer_global_rank, self.group) ) - ops.append(dist.P2POp(dist.isend, grad, peer_global_rank, self.group)) else: if grad is not None or grad_recv_stage is not None: - raise RuntimeError( + raise PipeliningMetadataError( f"[{self.stage_index}] for chunk {bwd_chunk_id} has gradients {grad} " f"and is expecting to send gradients to stage {grad_recv_stage}" ) @@ -523,34 +563,90 @@ def clear_runtime_states(self) -> None: # don't want such accumulation. for recv_tuple in self.args_recv_info.values(): # iterate over all chunks for a in recv_tuple: # iterate over all input args - if isinstance(a, _RecvInfo): + if not a.is_root_arg and a.buffer is not None: # Set to None is the newer and recommended way to clear grads, compared to `zero_()`. # See https://github.com/pytorch/pytorch/pull/92731 a.buffer.grad = None def _map_tensor_from_recv_info( self, - recv_infos: tuple[InputInfo, ...], + recv_infos: tuple[_RecvInfo, ...], ): """ Map tensors from recv infos to a list. """ def get_recv_tensor(info): - if isinstance(info, _RecvInfo): - return info.buffer - else: - raise AssertionError(f"Expected _RecvInfo but got {type(info)}") + if info.is_root_arg: + raise PipeliningMetadataError("Cannot get recv tensor from root arg") + return info.buffer return map_aggregate(cast(Argument, recv_infos), get_recv_tensor) - def _retrieve_recv_activations(self, fwd_chunk_id: int): + def _retrieve_recv_activations( + self, + fwd_chunk_id: int, + ): """ Retrieve the activations received for the current stage during forward. + Reconstructs DTensors if the inputs were DTensors. + Also validates DTensor metadata against expected values. """ recv_infos = self.args_recv_info[fwd_chunk_id] - activations = self._map_tensor_from_recv_info(recv_infos) - return activations + + activations = [] + for i, info in enumerate(recv_infos): + if not info.is_root_arg: + # Non-root args have valid buffer and tensor_meta + if info.buffer is None or info.tensor_meta is None: + raise PipeliningMetadataError( + f"Non-root arg '{info.input_name}' has None buffer or tensor_meta" + ) + # Effective requires_grad: metadata captures what the model + # produced, but the runtime context (has_backward, grad mode) + # determines whether we actually need gradients. + effective_requires_grad = ( + info.tensor_meta.requires_grad + and self.has_backward + and torch.is_grad_enabled() + ) + if isinstance(info.tensor_meta, _DTensorMeta): + # Buffer must not require grad so from_local stays out + # of the autograd graph (no grad_placements needed). + if info.buffer.requires_grad: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: recv buffer " + f"'{info.input_name}' unexpectedly requires grad " + f"before DTensor reconstruction" + ) + mesh = self._mesh_cache.get_mesh(info.tensor_meta.mesh_cache_key) + activation = DTensor.from_local( + info.buffer, + device_mesh=mesh, + placements=info.tensor_meta.placements, + shape=info.tensor_meta.global_shape, + stride=info.tensor_meta.global_stride, + run_check=False, + ).requires_grad_(effective_requires_grad) + else: + activation = info.buffer.requires_grad_(effective_requires_grad) + # Activation must be a leaf so backward terminates here. + if effective_requires_grad and not activation.is_leaf: + warnings.warn( + f"Stage {self.stage_index}: activation " + f"'{info.input_name}' is not a leaf " + f"(grad_fn={activation.grad_fn}); using " + f"retain_grad() as fallback", + stacklevel=2, + ) + activation.retain_grad() + activations.append(activation) + else: + raise PipeliningMetadataError( + f"_retrieve_recv_activations expected non-root _RecvInfo but got root arg at index {i}" + ) + + return tuple(activations) def _retrieve_recv_grads( self, @@ -558,10 +654,50 @@ def _retrieve_recv_grads( ): """ Retrieve the gradients received for the current stage during backward. + + Handles None gradients gracefully (for inputs that don't require grad). """ recv_infos = self.grad_recv_info[bwd_chunk_id] - grads = self._map_tensor_from_recv_info(recv_infos) - return grads + + grads: list[torch.Tensor | None] = [] + for i, info in enumerate(recv_infos): + if not isinstance(info, _RecvInfo): + raise PipeliningMetadataError( + f"Expected _RecvInfo but got {type(info)}" + ) + if not info.is_root_arg: + # Gradients can be None for non-differentiable outputs + if info.buffer is None: + if info.tensor_meta is not None: + raise PipeliningMetadataError( + f"Grad recv '{info.input_name}': buffer is None but tensor_meta is not None" + ) + grads.append(None) + continue + if info.tensor_meta is None: + raise PipeliningMetadataError( + f"Grad recv '{info.input_name}': buffer is not None but tensor_meta is None" + ) + if isinstance(info.tensor_meta, _DTensorMeta): + # Reconstruct DTensor gradient from local tensor + metadata + mesh = self._mesh_cache.get_mesh(info.tensor_meta.mesh_cache_key) + grad = DTensor.from_local( + info.buffer, + device_mesh=mesh, + placements=info.tensor_meta.placements, + shape=info.tensor_meta.global_shape, + stride=info.tensor_meta.global_stride, + run_check=False, + ) + else: + grad = info.buffer + grads.append(grad) + else: + raise PipeliningMetadataError( + f"grad_recv_info should not contain root args, but found one at index {i}" + ) + + return tuple(grads) def forward_maybe_with_nosync(self, *args, **kwargs): # If submod is wrapped with DDP, we use the `no_sync` context manager to @@ -691,7 +827,12 @@ def forward_one_chunk( composite_kwargs = kwargs or {} - self._validate_fwd_input(args, kwargs) + if self._runtime_validate: + self._validate_stage_tensors( + f"Stage {self.stage_index} forward inputs", + self._stage_meta.inputs, + composite_args, + ) # Compute forward try: @@ -727,7 +868,14 @@ def forward_one_chunk( fwd_chunk_id, map_debug_info(output), ) - self._validate_fwd_outputs(output_tuple) + # Validate outputs before P2P send; skipped for last stage (outputs + # go to loss/user, not via send/recv). + if self._runtime_validate and not self.is_last: + self._validate_stage_tensors( + f"Stage {self.stage_index} forward outputs", + self._stage_meta.outputs, + output_tuple, + ) # We return the original user-provided output, not normalized to tuple. # See [Note: pipeline model output type] @@ -776,6 +924,13 @@ def backward_one_chunk( else: # Otherwise, receive gradients from next stage grads_output = self._retrieve_recv_grads(bwd_chunk_id) + if self._runtime_validate: + # Validate backward input (output gradients) for DTensor metadata + self._validate_stage_tensors( + f"Stage {self.stage_index} backward input (output_grads)", + self._stage_meta.output_grads, + grads_output, + ) # If an input to the pipeline requires gradient, # `torch.autograd.backward` will accumulate the gradient into the # `.grad` field of such input @@ -828,8 +983,18 @@ def backward_one_chunk( ) # Save a placeholder for the dw_runner self.dw_runner[bwd_chunk_id] = lambda: None - - self.bwd_cache[bwd_chunk_id] = grads_input + # Note: grads_input may contain gradients for both args and kwargs (from fwd_cache), + # Kwargs are local to each stage and don't need gradient transmission. + # Validate backward output (input gradients) for DTensor metadata + assert self._stage_meta.inputs is not None # noqa: S101 + num_fwd_args = len(self._stage_meta.inputs) + if self._runtime_validate and not self.is_first: + self._validate_stage_tensors( + f"Stage {self.stage_index} backward output (input_grads)", + self._stage_meta.input_grads, + grads_input[:num_fwd_args], + ) + self.bwd_cache[bwd_chunk_id] = grads_input[:num_fwd_args] if self.is_last and not self.is_first: # Autograd dependencies: @@ -887,48 +1052,6 @@ def backward_weight_one_chunk(self, bwd_chunk_id: int, last_backward=False): "full", bwd_kwargs, last_backward=last_backward ) - def _validate_fwd_input(self, args, kwargs): - """Raises a RuntimeError if shapes of input args/kwargs do not match the shapes configured for this stage.""" - - if self.is_first: - # TODO why is there a separate recv_info for each pipeline chunk? - # kwen2501: to avoid passing a `fwd_chunk_id` to this function, we - # check all chunks against args_recv_info[0] - expected_args = self.args_recv_info[0] - else: - # We don't check inputs for non-0 stages assuming they don't accept - # user inputs in canonical pipeline scenarios - return - - if len(kwargs): - # TODO- need a mapping of kwarg to position in self.args_recv_info - # Without it, we are not 100% sure how to match the args and - # expected_args. - return - - # TODO- need a mapping of kwarg to position in self.args_recv_info - # maybe it's impossible to tell whether the len mismatches because - # (a) the user passed an extra arg or missed an arg - # (b) the user did not pass a kwarg, which has a default value baked into expected_args - expected_tensors_meta = [ - e.meta if isinstance(e, _RootArgPlaceholder) else e.buffer - for e in expected_args - ] - validate_tensors_metadata( - f"Stage {self.stage_index} forward inputs", expected_tensors_meta, args - ) - - def _validate_fwd_outputs(self, outputs: tuple[torch.Tensor, ...]): - """Raises a RuntimeError if this stage produces an output of unexpected shape/dtype. - Most likely, this could be cause either by incorrect user specification of output shapes, or because - shape inference was done on the original model but then at runtime the model is wrapped with something like - mixed precision which changes output dtype. - """ - expected_tensors_meta = self.get_outputs_meta() - validate_tensors_metadata( - f"Stage {self.stage_index} forward outputs", expected_tensors_meta, outputs - ) - def _get_init_p2p_neighbors_ops(self) -> list[dist.P2POp]: """ Get the operations to initialize the p2p communicators between previous and next stages. @@ -1054,7 +1177,7 @@ def __init__( node for node in pipe_info.graph.nodes if node.op == "call_module" ] if len(submod_nodes) != self.num_stages: - raise AssertionError( + raise PipeliningMetadataError( f"Number of submodules in pipe graph {len(submod_nodes)} does not match number of stages {self.num_stages}" ) @@ -1092,21 +1215,85 @@ def _move_submod_to_device(self): def _prepare_forward_infra( self, num_microbatches: int, - args: tuple[Any, ...], + args: tuple[Any, ...] | _StageForwardMeta | None, kwargs: dict[str, Any] | None = None, - ) -> tuple[Any, ...]: + has_backward: bool = False, + ) -> _StageForwardMeta | None: """ - Create send/recv infrastructures for activations (during forward) + Prepare forward infrastructure for traced pipeline. + + Metadata is created directly from graph placeholders with correct + ``requires_grad`` — received activations get ``requires_grad=True`` + when ``has_backward`` is set, fixing the fact that ``torch.export`` + traces under ``no_grad()``. + + ``_stage_meta.inputs`` is derived from recv infos and aligned with + ``forward_one_chunk``'s ``composite_args``: positional root inputs + on the first stage, received activations only on subsequent stages. """ - # TODO(whc) - # this method should be deleted once lazy buffer allocation is implemented - # for now, it ignores args/kwargs because it should not need to do shape inference + # Step 1: Create recv info for each microbatch. + # _create_act_recv_info is self-contained: it creates _TensorMeta + # directly from graph placeholder values with correct requires_grad. for chunk in range(num_microbatches): self.args_recv_info[chunk] = self._create_act_recv_info() - # Send info during forward for each activation + # Step 2: Derive _stage_meta.inputs from recv infos. + # forward_one_chunk builds composite_args as: + # - First stage: args (positional root inputs, excludes kwargs) + # - Non-first stages: received activations only (no root kwargs) + # _stage_meta.inputs must match composite_args for validation. + recv_infos = self.args_recv_info[0] + if self.is_first: + # All placeholders are root args. Only the first len(args) + # correspond to positional inputs (composite_args); the rest + # are kwargs passed separately via composite_kwargs. + # First stage always receives real tensor args, never _StageForwardMeta. + if not isinstance(args, tuple): + raise AssertionError("First stage requires real tensor args") + n_positional = len(args) + self._stage_meta.inputs = tuple( + info.tensor_meta # type: ignore[misc] + for info in recv_infos[:n_positional] + ) + else: + self._stage_meta.inputs = tuple( + info.tensor_meta # type: ignore[misc] + for info in recv_infos + if not info.is_root_arg + ) + + # Step 3: Create send info and output metadata. self.act_send_info = self._create_act_send_info() - return tuple() + + return None + + def _prepare_backward_infra( + self, + num_microbatches: int, + loss_fn: Callable[..., torch.Tensor] | None = None, + target: torch.Tensor | None = None, + received_grad_meta: _StageBackwardMeta | None = None, + ) -> _StageBackwardMeta | None: + """ + Prepare backward infrastructure for traced pipeline. + Derives input_grads metadata from inputs (plain tensors only). + + Note: DTensors are NOT supported in the traced frontend. + """ + # Derive input_grads from inputs (for plain tensors, grad shape == input shape) + if self._stage_meta.inputs is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: inputs metadata required for backward inference." + ) + + self._stage_meta.input_grads = _derive_grad_metas(self._stage_meta.inputs) + + # Setup backward recv info (calls _create_grad_recv_info which sets output_grads). + # Note: grad_send_info is created lazily in get_bwd_send_ops() since + # it mirrors args_recv_info (already populated during forward). + self._setup_backward_recv_info(num_microbatches) + + return None def get_stage_index_of_submod( self, @@ -1116,7 +1303,7 @@ def get_stage_index_of_submod( Given a submodule name, return the stage index of the submodule. """ if submod_name not in self.submod_to_stage_index: - raise AssertionError(f"Stage id of {submod_name} not found") + raise PipeliningMetadataError(f"Stage id of {submod_name} not found") return self.submod_to_stage_index[submod_name] @@ -1125,38 +1312,65 @@ def _create_act_recv_info( ): """ Create a tuple of `_RecvInfo` for inputs to the stage. + + Self-contained: creates ``_TensorMeta`` directly from graph + placeholder values with correct ``requires_grad``. + ``torch.export`` traces under ``no_grad()`` so traced metadata + always has ``requires_grad=False``; for received activations we + set ``requires_grad=True`` when ``has_backward`` is set. + + Note: DTensors are NOT supported in the traced frontend. """ def create_recv_tensor(placeholder, arg_node): - """ - Create a receive buffer for a placeholder. - """ example_value = placeholder.meta["val"] + + # Reject DTensors in traced frontend + if isinstance(example_value, DTensor): + raise PipeliningMetadataError( + f"{self.log_prefix} DTensor detected in traced pipeline input " + f"'{placeholder.name}'. DTensor metadata propagation is NOT " + f"supported for the traced frontend (_PipelineStage). " + f"Use the manual PipelineStage frontend for full DTensor support." + ) + if arg_node.op == "placeholder": - # This is a root level placeholder, thus an input argument to the entire model. - # We are likely at stage 0, hence no need to create a receive buffer. - return _RootArgPlaceholder(example_value) + # Root-level placeholder: an input argument to the entire + # model. Keep original metadata from the trace. + return _RecvInfo( + input_name=f"root_input_{placeholder.name}", + source=None, + buffer=None, + tensor_meta=_TensorMeta.from_tensor(example_value), + is_root_arg=True, + ) - # Figure out the source stage of this input + # Received activation from a previous stage. while arg_node.target is operator.getitem: - # If the input is a getitem, we need to go deeper arg_node = arg_node.args[0] if arg_node.op != "call_module": - raise AssertionError(f"Expecting call_module, got {arg_node.op}") + raise PipeliningMetadataError( + f"Expecting call_module, got {arg_node.op}" + ) src_stage = self.get_stage_index_of_submod(arg_node.name) - # Create a receive buffer for this placeholder + # Create metadata directly with correct requires_grad. + tensor_meta = _TensorMeta( + shape=example_value.shape, + stride=example_value.stride(), + dtype=example_value.dtype, + requires_grad=self.has_backward, + ) + logger.debug( "%s Creating recv buffer for input '%s' : %s, %s", self.log_prefix, placeholder.name, - example_value.shape, - example_value.dtype, + tensor_meta.shape, + tensor_meta.dtype, ) - buffer = _make_tensor_from_meta(example_value, self.device) - # In case there is backward pass, set requires_grad for receive buffers - # before first forward + buffer = _make_tensor_from_meta(tensor_meta, self.device) if self.has_backward: buffer.requires_grad_(True) @@ -1164,10 +1378,10 @@ def create_recv_tensor(placeholder, arg_node): arg_node.name, src_stage, buffer, + tensor_meta, ) - args_recv_info: list[InputInfo] = [] - # Filter out placeholder nodes from `self.submod` (a GraphModule) + args_recv_info: list[_RecvInfo] = [] placeholders = filter( # type: ignore[var-annotated] lambda node: node.op == "placeholder", # type: ignore[arg-type] self.submod.graph.nodes, # type: ignore[arg-type,union-attr] @@ -1175,15 +1389,12 @@ def create_recv_tensor(placeholder, arg_node): # `placeholders` are nodes internal to submod. # `self.node.args` are dependency nodes in the outer graph. # The two are 1:1. - for placeholder, arg_node in zip(placeholders, self.node.args): - # Create a receive buffer for this placeholder - recv_info = create_recv_tensor(placeholder, arg_node) - args_recv_info.append(recv_info) + for placeholder, arg_node in zip(placeholders, self.node.args, strict=True): + args_recv_info.append(create_recv_tensor(placeholder, arg_node)) logger.debug( "%s Activation recv / args info: %s", self.log_prefix, args_recv_info ) - # `args` is a Tuple, hence we will return a Tuple[InputInfo] return tuple(args_recv_info) def find_dst_rank( @@ -1207,14 +1418,12 @@ def find_dst_rank( def _create_act_send_info(self): """ - Create a dict of send info for activations. - The dict is of the form: - { - output_index: [dst_rank_0, dst_rank_1, ...], - ... - } - where the list of `dst_rank`s covers the case where an output value may - be consumed by multiple stages. + Create a dict of send info for activations and output metadata. + + Output metadata is created directly with correct ``requires_grad`` + (``torch.export`` traces under ``no_grad()``, so traced values + always have ``requires_grad=False``; at runtime, stage outputs + carry ``requires_grad=True`` when training). """ # Output index: List of receiver ranks act_send_info: dict[int, list] = {} @@ -1222,16 +1431,13 @@ def _create_act_send_info(self): for user in self.node.users: if user.target is operator.getitem: - # Recursively find the real destination gi_dsts = act_send_info.setdefault(out_idx, []) for gi_user in user.users: dst_rank = self.find_dst_rank(gi_user) if dst_rank is not None: gi_dsts.append(dst_rank) - # Next `getitem` will point to the next output index out_idx += 1 else: - # In case of single output value, `out_idx` will not increase dsts = act_send_info.setdefault(out_idx, []) dst_rank = self.find_dst_rank(user) if dst_rank is not None: @@ -1241,7 +1447,25 @@ def _create_act_send_info(self): output_vals: tuple[torch.Tensor] = tuple( v.meta["val"] for v in flatten_args(output_node.args) ) - self._configure_outputs_meta(output_vals) + # Reject DTensors and create output metadata directly with + # correct requires_grad. + output_metas: list[_TensorMeta] = [] + for i, val in enumerate(output_vals): + if isinstance(val, DTensor): + raise PipeliningMetadataError( + f"{self.log_prefix} DTensor detected in traced pipeline output index {i}. " + f"DTensor metadata propagation is NOT supported for the traced frontend " + f"(_PipelineStage). Use the manual PipelineStage frontend for full DTensor support." + ) + output_metas.append( + _TensorMeta( + shape=val.shape, + stride=val.stride(), + dtype=val.dtype, + requires_grad=self.has_backward, + ) + ) + self._stage_meta.outputs = tuple(output_metas) logger.debug("%s Send info: %s", self.log_prefix, act_send_info) return act_send_info @@ -1249,50 +1473,79 @@ def _create_act_send_info(self): def _get_output_node(self): output_nodes = [node for node in self.submod.graph.nodes if node.op == "output"] # type: ignore[union-attr] if len(output_nodes) != 1: - raise AssertionError(f"Expected 1 output node, got {len(output_nodes)}") + raise PipeliningMetadataError( + f"Expected 1 output node, got {len(output_nodes)}" + ) output_node = output_nodes[0] return output_node - def _create_grad_recv_info( - self, - act_send_info: dict, - ) -> tuple[_RecvInfo, ...]: + def _create_grad_recv_info(self, act_send_info: dict) -> tuple[_RecvInfo, ...]: """ Create a tuple of `_RecvInfo` for gradients. + Reuses output metadata from _stage_meta.outputs (populated by _create_act_send_info). """ - # Dict[output_index, _RecvInfo] - grad_recv_info: dict[int, _RecvInfo] = {} - output_node = self._get_output_node() + if self._stage_meta.outputs is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: outputs metadata required for grad recv info. " + f"Ensure _create_act_send_info is called first." + ) - # The output node may take multiple args, meaning the submod having multiple output values. - output_vals = flatten_args(output_node.args) + outputs_meta = self._stage_meta.outputs + output_grads_metas: list[TensorMeta | None] = [] + grad_recv_infos: list[_RecvInfo] = [] - for out_idx, dst_list in act_send_info.items(): - if not dst_list: - # No actual receiver for activation so no grad coming back - continue + for out_idx, out_meta in enumerate(outputs_meta): + dst_list = act_send_info.get(out_idx, []) - output = output_vals[out_idx] - example_value = output.meta["val"] - logger.debug( - f"{self.log_prefix} Creating grad recv buffer for output {output.name} " # noqa: G004 - f": {example_value.shape}, {example_value.dtype}" - ) + # Determine the source stage for gradients + grad_src = dst_list[0] if dst_list else self.stage_index + 1 - # TODO: otherwise needs grad accumulation - if len(dst_list) != 1: - raise AssertionError("Backward of skip connections not supported yet") - grad_src = dst_list[0] - grad_recv_info[out_idx] = _RecvInfo( - f"{grad_src}", # noqa: G004 - grad_src, - _make_tensor_from_meta(example_value, self.device), - ) + # Check if this output needs gradients + if not dst_list or not out_meta.requires_grad: + output_grads_metas.append(None) + grad_recv_infos.append( + _RecvInfo( + input_name=f"recv_grad_for_{self.stage_index}_none_{out_idx}", + source=grad_src, + buffer=None, + tensor_meta=None, + ) + ) + else: + # Derive grad metadata from output metadata (same shape, requires_grad=False) + grad_meta = _TensorMeta( + shape=out_meta.shape, + stride=out_meta.stride, + dtype=out_meta.dtype, + requires_grad=False, + ) + output_grads_metas.append(grad_meta) - # Convert to tuple for convenience in get_ops and retrieve tensor - grad_recv_info_tuple = tuple(grad_recv_info.values()) - logger.debug("%s Grad recv info: %s", self.log_prefix, grad_recv_info_tuple) - return grad_recv_info_tuple + if len(dst_list) != 1: + raise PipeliningMetadataError( + "Backward of skip connections not supported yet" + ) + + logger.debug( + "%s Creating grad recv buffer for output %s : %s, %s", + self.log_prefix, + out_idx, + grad_meta.shape, + grad_meta.dtype, + ) + + grad_recv_infos.append( + _RecvInfo( + input_name=f"recv_grad_for_{self.stage_index}_from_{grad_src}", + source=grad_src, + buffer=_make_tensor_from_meta(grad_meta, self.device), + tensor_meta=grad_meta, + ) + ) + + self._stage_meta.output_grads = tuple(output_grads_metas) + logger.debug("%s Grad recv info: %s", self.log_prefix, grad_recv_infos) + return tuple(grad_recv_infos) # A helper function to create a pipeline stage based on traced pipeline information @@ -1327,26 +1580,33 @@ def build_stage( class PipelineStage(_PipelineStageBase): - """ - A class representing a pipeline stage in a pipeline parallelism setup. + """A pipeline stage for pipeline parallelism with sequential model partitioning. + + Supports both **static** and **dynamic** metadata inference: - PipelineStage assumes sequential partitioning of the model, i.e. the model is split into chunks where outputs from - one chunk feed into inputs of the next chunk, with no skip connections. + Static mode: + All of ``input_args``, ``output_args`` (and ``input_grads``/``output_grads`` + when DTensors are present) are provided at construction time. - PipelineStage performs runtime shape/dtype inference automatically by propagating the outputs from stage0 to - stage1 and so forth, in linear order. To bypass shape inference, pass the `input_args` and `output_args` to each - PipelineStage instance. + Dynamic mode: + Metadata is inferred from the first microbatch at runtime; any + statically provided args are used for validation only. Args: - submodule (nn.Module): The PyTorch module wrapped by this stage. - stage_index (int): The ID of this stage. - num_stages (int): The total number of stages. - device (torch.device): The device where this stage is located. - input_args (Union[torch.Tensor, Tuple[torch.tensor]], optional): The input arguments for the submodule. - output_args (Union[torch.Tensor, Tuple[torch.tensor]], optional): The output arguments for the submodule. - group (dist.ProcessGroup, optional): The process group for distributed training. If None, default group. - dw_builder (Optional[Callable[[], Callable[..., None]]): If provided, dw_builder will build a new dw_runner function - that will the W action (input weights) for F, I, W (Fwd, Input, Weight) zero bubble schedules. + submodule: The ``nn.Module`` wrapped by this stage. + stage_index: Zero-based stage ID. + num_stages: Total number of stages in the pipeline. + device: Device this stage runs on. + input_args: Example input tensors (single tensor or tuple). Optional. + output_args: Example output tensors. Optional. + output_grads: Example output gradients (received from next stage). Optional. + input_grads: Example input gradients (sent to previous stage). Optional. + group: Process group for P2P communication. Defaults to the + world process group. + dw_builder: Builder for deferred weight-update runners used by + zero-bubble (F/I/W) schedules. + get_mesh: `GetMeshCallback` used during + dynamic DTensor inference. Ignored in fully static DTensor mode. """ def __init__( @@ -1357,247 +1617,633 @@ def __init__( device: torch.device, input_args: torch.Tensor | tuple[torch.Tensor, ...] | None = None, output_args: torch.Tensor | tuple[torch.Tensor, ...] | None = None, + output_grads: torch.Tensor | tuple[torch.Tensor | None, ...] | None = None, + input_grads: torch.Tensor | tuple[torch.Tensor | None, ...] | None = None, group: dist.ProcessGroup | None = None, dw_builder: Callable[[], Callable[..., None]] | None = None, + get_mesh: GetMeshCallback | None = None, ): super().__init__(submodule, stage_index, num_stages, device, group, dw_builder) - self.inputs: list[torch.Tensor] | None = None - self.inputs_meta: tuple[torch.Tensor, ...] | None = None - # Note: inputs and submod should ideally be on meta device. We decided not to assert this (yet) because it - # might be breaking for existing users. - if input_args is None: - if output_args is not None: - raise AssertionError( - "If specifying output_args, input_args must also be specified. " - "Otherwise, shape inference will be performed at runtime" - ) - else: - self.inputs_meta = ( - (input_args,) if isinstance(input_args, torch.Tensor) else input_args - ) - if output_args is None: - logger.warning( - "Deprecation warning: passing input_args and performing init-time shape inference is deprecated. " - "PipelineStage now supports runtime shape inference using the real inputs provided to schedule step(). " - "Either delete `input_args` arg to `PipelineStage` to opt-into runtime shape inference, " - "or additionally pass `output_args` to `PipelineStage` to fully override shape inference. " + + self._mesh_cache = _MeshCache(get_mesh_cb=get_mesh) + self._inference_mode: InferenceMode | None = None + self._fwd_outputs_for_bwd_meta: tuple[torch.Tensor, ...] | None = None + self._fwd_inputs_for_bwd_meta: tuple[torch.Tensor, ...] | None = None + self._fwd_kwargs_tensors_for_bwd_meta: tuple[torch.Tensor, ...] | None = None + + # Validate and normalize args to tuples + inputs = validate_and_normalize_to_tuple(input_args) + outputs = validate_and_normalize_to_tuple(output_args) + in_grads = validate_and_normalize_to_tuple(input_grads, allow_none=True) + out_grads = validate_and_normalize_to_tuple(output_grads, allow_none=True) + + self._user_meta = _StageMeta( + inputs=extract_tensor_metas(inputs), + outputs=extract_tensor_metas(outputs), + input_grads=extract_tensor_metas(in_grads, allow_none=True), + output_grads=extract_tensor_metas(out_grads, allow_none=True), + ) + + # Cache meshes from user-provided DTensors + for args in (inputs, outputs, in_grads, out_grads): + if args is not None: + self._mesh_cache.update_from_tensors(args) + + # Validate DTensor↔grad correspondence independently for inputs and outputs + if self._user_meta.has_dtensors(): + if inputs and in_grads: + validate_static_arg_grad_correspondence( + self.stage_index, inputs, in_grads, is_input=True ) - try: - with torch.no_grad(): - output_args = submodule(*self.inputs_meta) - output_args = tree_map_only( - torch.Tensor, lambda x: x.to("meta"), output_args - ) - except Exception as e: - raise RuntimeError( - "Failed to perform pipeline shape inference- are your inputs on the same device as your module?" - ) from e - if output_args is None: - raise AssertionError( - "If passing input_args, also pass output_args to override shape inference" + if outputs and out_grads: + validate_static_arg_grad_correspondence( + self.stage_index, outputs, out_grads, is_input=False ) - self._configure_outputs_meta( - (output_args,) if isinstance(output_args, torch.Tensor) else output_args + + def _recv_meta(self, src_stage: int) -> Any: + """Receive metadata object from a stage on a different rank via P2P.""" + objects: list[Any] = [None] + dist.recv_object_list( + objects, + src=self._resolve_peer_global_rank(src_stage), + group=self.group, + device=self.device, + use_batch=True, + ) + if len(objects) != 1: + raise PipeliningMetadataError( + f"Expected exactly one object to be received but got: {len(objects)}" ) + return objects[0] + + def _send_meta(self, meta: Any, dst_stage: int) -> None: + """Send metadata object to a stage on a different rank via P2P.""" + dist.send_object_list( + [meta], + dst=self._resolve_peer_global_rank(dst_stage), + group=self.group, + device=self.device, + use_batch=True, + ) + + def _is_same_rank(self, other_stage: int) -> bool: + """Check if another stage is on the same rank as this stage.""" + return self.stage_index_to_group_rank[other_stage] == self.group_rank - # these are the buffers used in backwards send/recv, they are allocated later - self.outputs_grad: list[torch.Tensor] = [] + def _warmup_forward_vote( + self, has_backward: bool, received_acc: torch.Tensor | None = None + ) -> torch.Tensor: + """Forward phase of the warm-up vote protocol (stage 0 → N−1). + + Each stage computes a vote (1 = STATIC, 0 = DYNAMIC) based on + ``InferenceMode.needs_dynamic``, multiplies it with the accumulated + product from the previous stage, and forwards the result to the next + stage. The final product at stage N−1 is 1 iff *every* stage voted + STATIC. + + Args: + has_backward: Whether the schedule includes a backward pass. + received_acc: Accumulated product tensor from the previous + same-rank stage (V-schedule), or ``None`` for the first + stage / cross-rank. + + Returns: + The accumulated product tensor after this stage's vote. + """ + my_vote = 0 if InferenceMode.needs_dynamic(self._user_meta, has_backward) else 1 + + my_vote_t = torch.tensor([my_vote], dtype=torch.int32, device=self.device) + + if self.is_first: + acc = my_vote_t + elif self._is_same_rank(self.stage_index - 1): + assert received_acc is not None # noqa: S101 + acc = received_acc * my_vote_t + else: + peer_global = self._resolve_peer_global_rank(self.stage_index - 1) + acc = torch.zeros(1, dtype=torch.int32, device=self.device) + dist.recv(acc, src=peer_global, group=self.group) + acc = acc * my_vote_t - dbg_str = ( - f"Finished pipeline stage init, {self.stage_index=}, {self.is_first=}, " # noqa: G004 - f"{self.is_last=}, {self.num_stages=}, " + if not self.is_last and not self._is_same_rank(self.stage_index + 1): + peer_global = self._resolve_peer_global_rank(self.stage_index + 1) + dist.send(acc, dst=peer_global, group=self.group) + + return acc + + def _warmup_backward_result( + self, received_result: torch.Tensor | None = None + ) -> torch.Tensor: + """Backward phase of the warm-up vote protocol (stage N−1 → 0). + + Propagates the final accumulated product (computed in the forward + phase) back through the pipeline so every stage learns the global + inference mode. + + Args: + received_result: Result tensor from the next same-rank stage + (V-schedule), or ``None`` for the last stage / cross-rank. + + Returns: + The global vote result tensor for this stage. + """ + if self.is_last or self._is_same_rank(self.stage_index + 1): + assert received_result is not None # noqa: S101 + result = received_result + else: + peer_global = self._resolve_peer_global_rank(self.stage_index + 1) + result = torch.zeros(1, dtype=torch.int32, device=self.device) + dist.recv(result, src=peer_global, group=self.group) + + if not self.is_first and not self._is_same_rank(self.stage_index - 1): + peer_global = self._resolve_peer_global_rank(self.stage_index - 1) + dist.send(result, dst=peer_global, group=self.group) + + return result + + def _compute_outputs( + self, + *args: torch.Tensor, + module: torch.nn.Module, + **kwargs: Any, + ) -> torch.Tensor | tuple[torch.Tensor, ...] | list[torch.Tensor] | None: + """Compute outputs of the submodule.""" + return module(*args, **kwargs) + + def _compute_input_grads( + self, + outputs: list[torch.Tensor], + all_fwd_inputs: list[torch.Tensor], + grad_outputs: list[torch.Tensor | None] | None = None, + ) -> tuple[torch.Tensor | None, ...]: + """Compute input gradients via :func:`_autograd_grad_for_inputs`.""" + return _autograd_grad_for_inputs( + outputs, + all_fwd_inputs, + grad_outputs, ) - if self.inputs_meta is not None: - dbg_str += ( - f"inputs: {[inp.shape for inp in self.inputs_meta]}, " - f"output: {[output.shape for output in self.get_outputs_meta()]}" - ) + + def _to_tensor(self, arg: torch.Tensor | TensorMeta) -> torch.Tensor: + """Convert a tensor or metadata to a real tensor on ``self.device``. + + Real tensors are detached and re-set requires_grad to create a fresh + autograd leaf, isolating metadata inference from the user's graph. + TensorMeta is materialized as an empty tensor (or DTensor via mesh cache). + """ + if isinstance(arg, torch.Tensor): + return arg.detach().requires_grad_(arg.requires_grad) + elif isinstance(arg, TensorMeta): + if isinstance(arg, _DTensorMeta): + mesh = self._mesh_cache.get_mesh(arg.mesh_cache_key) + return arg.to_dtensor(self.device, mesh) + else: + return arg.to_tensor(self.device) else: - dbg_str += " running shape-inference at runtime" + raise PipeliningMetadataError( + f"Unsupported type {type(arg)} for _to_tensor: {arg}" + ) - logger.debug(dbg_str) + def _ones_from_metadata(self, meta: TensorMeta) -> torch.Tensor: + """Create a ones tensor from metadata for backward inference grad_outputs.""" + local_ones = torch.ones( + meta.shape, + dtype=meta.dtype, + device=self.device, + ) + if isinstance(meta, _DTensorMeta): + mesh = self._mesh_cache.get_mesh(meta.mesh_cache_key) + return DTensor.from_local( + local_ones, + device_mesh=mesh, + placements=meta.placements, + shape=meta.global_shape, + stride=meta.global_stride, + run_check=False, + ) + return local_ones - def _shape_inference( + def _forward_metadata_inference( self, - args: tuple[Any, ...], + args: tuple[torch.Tensor, ...] | _StageForwardMeta | None, kwargs: dict[str, Any] | None = None, - ): - if kwargs is None: - kwargs = {} - if args is None: - raise AssertionError("Args may be an empty tuple but not None") - - # We skip recv communication if we're the first stage, but also if the previous stage is on the same rank - # and can pass its output shapes in as args instead of using send/recv. - if ( - self.is_first - # if not first stage, then check if prev stage is on the same rank - or self.stage_index_to_group_rank[self.stage_index - 1] == self.group_rank - ): - logger.debug( - "Shape inference: stage %s skipping recv, because shape info passed in via `args`", - self.stage_index, - ) - args = tree_map_only(torch.Tensor, lambda x: x.to("meta"), args) + has_backward: bool = False, + ) -> _StageForwardMeta | None: + """Run forward metadata inference (Stage 0 → N). + + Args: + args: Real tensors (first stage), ``_StageForwardMeta`` + (same-rank), or ``None`` (cross-rank P2P). + kwargs: Keyword arguments forwarded to the submodule. + has_backward: Whether backward inference follows. + + Returns: + ``_StageForwardMeta`` for the next stage, or ``None`` if sent via P2P. + """ + kwargs = kwargs or {} + + # === RECEIVE: Get input metadata and create meta tensors === + if self.is_first: + # First stage: extract metadata from real tensors + if args is None or isinstance(args, _StageForwardMeta): + raise PipeliningMetadataError( + f"Stage {self.stage_index}: First stage requires real tensors, " + f"got {type(args).__name__}." + ) + tensor_args = validate_and_normalize_to_tuple(args) + assert tensor_args is not None # noqa: S101 + self._stage_meta.inputs = extract_tensor_metas(tensor_args) + inference_args = tuple(self._to_tensor(a) for a in tensor_args) + elif self._is_same_rank(self.stage_index - 1): + # Same-rank: _StageForwardMeta passed via argument + if not isinstance(args, _StageForwardMeta): + raise PipeliningMetadataError( + f"Stage {self.stage_index}: Expected _StageForwardMeta from same-rank " + f"previous stage, got {type(args).__name__}." + ) + self._stage_meta.inputs = args.forward_metas + inference_args = tuple(self._to_tensor(m) for m in args.forward_metas) else: - if len(args) != 0: - raise AssertionError( - "Can't supply input args for shape inference on non-first stage" + # Cross-rank: receive _StageForwardMeta via P2P + recv_meta = self._recv_meta(self.stage_index - 1) + if not isinstance(recv_meta, _StageForwardMeta): + raise PipeliningMetadataError( + f"Stage {self.stage_index}: Expected _StageForwardMeta from P2P, " + f"got {type(recv_meta).__name__}." ) - objects = [None] - logger.debug( - "Shape inference: stage %s receiving from stage %s", - self.stage_index, - self.stage_index - 1, - ) - dist.recv_object_list( - objects, - src=dist.get_global_rank( - self.group or dist.distributed_c10d._get_default_group(), - self.stage_index_to_group_rank[self.stage_index - 1], - ), - group=self.group, - device=self.device, - use_batch=True, + self._stage_meta.inputs = recv_meta.forward_metas + inference_args = tuple(self._to_tensor(m) for m in recv_meta.forward_metas) + + inference_kwargs = { + k: self._to_tensor(v) if isinstance(v, torch.Tensor) else v + for k, v in kwargs.items() + } + + # Isolate metadata inference from user's grad context. + # has_backward → enable_grad() so backward tracing sees grad_fn; + # no backward → no_grad() for cross-rank consistency. + ctx = torch.enable_grad() if has_backward else torch.no_grad() + with ctx: + outputs = self._compute_outputs( + *inference_args, module=self.submod, **inference_kwargs ) - recv_args = objects[0] - if not isinstance(recv_args, tuple): - raise AssertionError(f"Expected tuple, got {type(recv_args)}") - args = recv_args - - # cache input shapes for use during recv buffer allocation - self.inputs_meta = args - args = tree_map_only( - torch.Tensor, lambda x: torch.zeros_like(x, device=self.device), args - ) - # set attributes needed for forward - with torch.no_grad(): - outputs = self.submod(*args, **kwargs) + # Normalize outputs to tuple + outputs = validate_and_normalize_to_tuple(outputs) - # if single tensor, convert so it is always a list - if isinstance(outputs, torch.Tensor): - outputs = [outputs] + self._stage_meta.outputs = extract_tensor_metas(outputs) - # communicate meta outputs not real outputs for two reasons - # 1 - its faster (esp. since obj coll pickles tensor data!) - # 2 - avoid activating a cuda context for the src rank when unpickling on the recv end! - outputs_meta = tuple( - tree_map_only(torch.Tensor, lambda x: x.to("meta"), outputs) - ) - logger.debug( - "Shape inference: stage %s inputs %s, outputs %s", - self.stage_index, - self.inputs_meta, - outputs_meta, + # Store for backward metadata inference (always, even during eval) + fwd_kwargs_tensors = tuple( + v for v in flatten_args(inference_kwargs) if isinstance(v, torch.Tensor) ) - self._configure_outputs_meta(outputs_meta) - - # Passing outputs to the next stage: - # two cases- - # 1. Usually: use send/recv communication to pass the output - # 2. Special case: for V-schedules, 2 'adjacent' stages (e.g. stage 3, 4 in an 8-stage 4-rank V) - # pass their shape info via return value and function args rather than send/recv. - if ( - self.is_last - # if not last stage, then check if next stage is on the same rank - or self.stage_index_to_group_rank[self.stage_index + 1] == self.group_rank - ): - # Case (2) above: pass shape info via return value and caller passes it as args to next stage's - # _shape_inference call - logger.debug( - "Shape inference: stage %s skipping send to next stage", - self.stage_index, + self._fwd_outputs_for_bwd_meta = outputs + self._fwd_inputs_for_bwd_meta = inference_args + self._fwd_kwargs_tensors_for_bwd_meta = fwd_kwargs_tensors + + # === SEND: Pass output metadata to next stage === + if self._stage_meta.outputs is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: output metadata is required for forward inference." ) + fwd_meta = _StageForwardMeta(forward_metas=self._stage_meta.outputs) + if self.is_last or self._is_same_rank(self.stage_index + 1): + # Same-rank or last: return for caller to pass + return fwd_meta else: - # Case (1): send shapes via send operation, and ensure not to return it to the caller - logger.debug( - "Shape inference: stage %s sending to stage %s", - self.stage_index, - self.stage_index + 1, + # Cross-rank: send via P2P + self._send_meta(fwd_meta, self.stage_index + 1) + return None + + def _backward_metadata_inference( + self, + loss_fn: Callable[..., torch.Tensor] | None = None, + target: torch.Tensor | None = None, + received_grad_meta: _StageBackwardMeta | None = None, + ) -> _StageBackwardMeta | None: + """Run backward metadata inference (Stage N → 0). + + Args: + loss_fn: Loss function (required for the last stage). + target: Target tensor (required for the last stage). + received_grad_meta: Grad metadata from next same-rank stage + (V-schedule only). + + Returns: + ``_StageBackwardMeta`` for the previous stage, or ``None`` if sent via P2P. + """ + fwd_outputs = self._fwd_outputs_for_bwd_meta + fwd_inputs = self._fwd_inputs_for_bwd_meta + if fwd_outputs is None or fwd_inputs is None: + raise PipeliningMetadataError( + "Backward metadata inference requires forward metadata inference to run first" ) - dist.send_object_list( - [outputs_meta], - dst=dist.get_global_rank( - self.group or dist.distributed_c10d._get_default_group(), - self.stage_index_to_group_rank[self.stage_index + 1], - ), - group=self.group, - device=self.device, - use_batch=True, + kwargs_tensors = self._fwd_kwargs_tensors_for_bwd_meta or () + all_fwd_inputs = list(fwd_inputs) + list(kwargs_tensors) + # Clear temporary storage early — local refs are sufficient from here + self._fwd_outputs_for_bwd_meta = None + self._fwd_inputs_for_bwd_meta = None + self._fwd_kwargs_tensors_for_bwd_meta = None + # === RECEIVE: Get output grad metadata (except last stage) === + if self.is_last: + if loss_fn is None or target is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: loss_fn and target required for last stage" + ) + inference_target = self._to_tensor(target) + loss = loss_fn( + fwd_outputs[0] if len(fwd_outputs) == 1 else fwd_outputs, + inference_target, + ) + self._stage_meta.output_grads = None + all_input_grads = self._compute_input_grads( + [loss], + all_fwd_inputs, ) - outputs_meta = tuple() + else: + # Non-last stage: receive grad metadata from next stage + if self._is_same_rank(self.stage_index + 1): + # Same-rank: _StageBackwardMeta passed via argument + if not isinstance(received_grad_meta, _StageBackwardMeta): + raise PipeliningMetadataError( + f"Stage {self.stage_index}: Expected _StageBackwardMeta from same-rank " + f"next stage, got {type(received_grad_meta).__name__}." + ) + self._stage_meta.output_grads = received_grad_meta.backward_metas + else: + # Cross-rank: receive _StageBackwardMeta via P2P + recv_meta = self._recv_meta(self.stage_index + 1) + if not isinstance(recv_meta, _StageBackwardMeta): + raise PipeliningMetadataError( + f"Stage {self.stage_index}: Expected _StageBackwardMeta from P2P, " + f"got {type(recv_meta).__name__}." + ) + self._stage_meta.output_grads = recv_meta.backward_metas + + # === COMPUTE: Build grad_outputs and compute input grads === + # Extract output tensors and corresponding grad_outputs from metadata + # Must iterate together to maintain alignment + if self._stage_meta.output_grads is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: output_grads metadata is required for backward inference." + ) + stage_output_grad_metas = self._stage_meta.output_grads + + filtered_fwd_outputs: list[torch.Tensor] = [] + filtered_output_grads: list[torch.Tensor | None] = [] + + for idx, (fwd_out, grad_meta) in enumerate( + zip(fwd_outputs, stage_output_grad_metas, strict=True) + ): + # Match _backward.py behavior: skip if output doesn't require grad AND has no grad_fn + if not fwd_out.requires_grad: + if grad_meta is not None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: output {idx} requires_grad=False, " + f"but output_grads metadata is provided: {grad_meta}." + ) + continue + filtered_fwd_outputs.append(fwd_out) + # For outputs that require grad, include them even if grad_meta is None + # (runtime passes None grad_outputs to autograd.backward in this case) + filtered_output_grads.append( + self._ones_from_metadata(grad_meta) if grad_meta else None + ) - return outputs_meta + if filtered_fwd_outputs: + all_input_grads = self._compute_input_grads( + filtered_fwd_outputs, all_fwd_inputs, filtered_output_grads + ) + # Free intermediate references early + filtered_fwd_outputs.clear() + filtered_output_grads.clear() + all_fwd_inputs.clear() + # Only positional input grads flow to previous stage + else: + all_input_grads = tuple(None for _ in range(len(all_fwd_inputs))) + + input_grads = all_input_grads[: len(fwd_inputs)] + self._stage_meta.input_grads = tuple( + extract_tensor_meta(g) if isinstance(g, torch.Tensor) else None + for g in input_grads + ) + + # === SEND: Pass input grad metadata to previous stage === + bwd_meta = _StageBackwardMeta(backward_metas=self._stage_meta.input_grads) + + if self.is_first or self._is_same_rank(self.stage_index - 1): + # First rank or Same-rank: return for caller to pass + return bwd_meta + else: + # Cross-rank: send via P2P + self._send_meta(bwd_meta, self.stage_index - 1) + return None + + def _post_metadata_inference_cleanup(self) -> None: + """Clean up FSDP side effects (unsharded params, stale grads, stored + tensors) after metadata inference with real tensors. + """ + # Clear stored inference tensors (frees autograd graph + activations) + self._fwd_outputs_for_bwd_meta = None + self._fwd_inputs_for_bwd_meta = None + self._fwd_kwargs_tensors_for_bwd_meta = None + + # Metadata inference runs real fwd/bwd, which unshards FSDP params and + # accumulates grads. Reshard to free memory and clear stale grads. + for module in self.submod.modules(): + if isinstance(module, FSDPModule): + module.reshard() + for param in module.parameters(): + param.grad = None + + def _prepare_backward_infra( + self, + num_microbatches: int, + loss_fn: Callable[..., torch.Tensor] | None = None, + target: torch.Tensor | None = None, + received_grad_meta: "_StageBackwardMeta | None" = None, + ) -> "_StageBackwardMeta | None": + """Run backward metadata inference and prepare backward infrastructure. + + Returns: + ``_StageBackwardMeta`` for the previous same-rank stage, or ``None``. + """ + grad_meta_result: _StageBackwardMeta | None = None + if self._inference_mode == InferenceMode.DYNAMIC: + # DYNAMIC mode: run backward metadata inference + # received_grad_meta is used for same-rank V-schedule stages + grad_meta_result = self._backward_metadata_inference( + loss_fn=loss_fn, + target=target, + received_grad_meta=received_grad_meta, + ) + # Validate dynamically inferred metadata against user-provided metadata + self._validate_inferred_metadata() + else: + # STATIC mode: metadata comes from user inputs, no validation needed + self._stage_meta.input_grads = self._user_meta.input_grads + self._stage_meta.output_grads = self._user_meta.output_grads + # For STATIC mode with plain tensors, if output_grads is not set but + # we have outputs, derive output_grads from outputs. + # (gradient shape == output shape, but requires_grad=False for gradients) + if self._stage_meta.output_grads is None: + if self._stage_meta.outputs is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: output metadata is required for backward inference." + ) + self._stage_meta.output_grads = _derive_grad_metas( + self._stage_meta.outputs + ) + # Similarly, derive input_grads from inputs if not provided + if self._stage_meta.input_grads is None: + if self._stage_meta.inputs is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: input metadata is required for backward inference." + ) + self._stage_meta.input_grads = _derive_grad_metas( + self._stage_meta.inputs + ) + + # Note: grad_send_info is created lazily in get_bwd_send_ops() since + # it mirrors args_recv_info (already populated during forward). + self._setup_backward_recv_info(num_microbatches) + return grad_meta_result + + def _validate_inferred_metadata(self) -> None: + """Validate dynamically inferred metadata against user-provided metadata.""" + pairs = [ + (self._user_meta.inputs, self._stage_meta.inputs, "input"), + (self._user_meta.outputs, self._stage_meta.outputs, "output"), + (self._user_meta.input_grads, self._stage_meta.input_grads, "input_grad"), + ( + self._user_meta.output_grads, + self._stage_meta.output_grads, + "output_grad", + ), + ] + for user_val, stage_val, label in pairs: + if user_val and stage_val: + validate_tensors_metadata( + f"Stage {self.stage_index} {label}", + user_val, + stage_val, + warn_on_mismatch=True, + ) def _prepare_forward_infra( self, num_microbatches: int, - args: tuple[Any, ...], + args: tuple[Any, ...] | _StageForwardMeta | None, kwargs: dict[str, Any] | None = None, - ) -> tuple[Any, ...]: - # TODO move self.device to an argument from step API (from its input tensors)? - if num_microbatches is None: - raise AssertionError("TODO fix num_microbatches") - - outputs: tuple[Any, ...] = tuple() - if self.inputs_meta is None: - outputs = self._shape_inference(args, kwargs) - - if self.inputs_meta is None: - raise AssertionError("Expected inputs_meta to be set after shape inference") - # Receive info during forward - # TODO: create args_recv_info lazily? (same needed for PipelineStage) + has_backward: bool = False, + ) -> _StageForwardMeta | None: + """Prepare the stage infrastructure for forward pass. + + Returns: + ``_StageForwardMeta`` for next stage (same-rank), or ``None`` if sent via P2P. + """ + if self._inference_mode is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: inference mode not set. " + f"Run warmup vote protocol first." + ) + + fwd_meta_output: _StageForwardMeta | None = None + + if self._inference_mode == InferenceMode.DYNAMIC: + # DYNAMIC mode: run forward metadata inference + # args may be _StageForwardMeta for same-rank V-schedule stages + fwd_meta_output = self._forward_metadata_inference( + args, kwargs, has_backward + ) + # Validate dynamically inferred metadata against user-provided metadata + self._validate_inferred_metadata() + # STATIC mode: metadata comes from user inputs, no validation needed + else: + self._stage_meta.inputs = self._user_meta.inputs + self._stage_meta.outputs = self._user_meta.outputs + + # Setup recv and send info + self._setup_forward_recv_info(num_microbatches, has_backward) + self._setup_forward_send_info() + + return fwd_meta_output + + def _setup_forward_recv_info( + self, num_microbatches: int, has_backward: bool + ) -> None: + """Setup receive info for forward pass.""" + if self._stage_meta.inputs is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: inputs metadata required for recv info." + ) for chunk_id in range(num_microbatches): - if not self.is_first: - # We assume that we always receive from stage - 1 - recv_infos = tuple( + if self.is_first: + # First stage: all inputs are root arguments (no recv needed) + self.args_recv_info[chunk_id] = tuple( _RecvInfo( - f"recv_for_{self.stage_index}_from_{self.stage_index - 1}", - self.stage_index - 1, - _make_tensor_from_meta(inp, self.device), + input_name=f"root_input_{idx}", + source=None, + buffer=None, + tensor_meta=meta, + is_root_arg=True, ) - for inp in self.inputs_meta + for idx, meta in enumerate(self._stage_meta.inputs) ) - # In case there is backward pass, set requires_grad for receive buffers - if self.has_backward: - for r in recv_infos: - r.buffer.requires_grad_(True) - - self.args_recv_info[chunk_id] = recv_infos else: + # Non-first stages: receive from previous stage self.args_recv_info[chunk_id] = tuple( - _RootArgPlaceholder(i) for i in self.inputs_meta + _RecvInfo( + input_name=f"recv_for_{self.stage_index}_from_{self.stage_index - 1}", + source=self.stage_index - 1, + buffer=_make_tensor_from_meta(meta, self.device), + tensor_meta=meta, + ) + for meta in self._stage_meta.inputs ) - # Send info during forward for each activation - # only need the rank that is being sent to + def _setup_forward_send_info(self) -> None: + """Setup send info for forward pass.""" self.act_send_info: dict[int, list] = {} - - for idx in range(len(self.get_outputs_meta())): - # We assume we always send to stage + 1 - if not self.is_last: - self.act_send_info[idx] = [self.stage_index + 1] - else: - self.act_send_info[idx] = [] - - return outputs + if self._stage_meta.outputs is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: outputs metadata required for recv info." + ) + for idx in range(len(self._stage_meta.outputs)): + self.act_send_info[idx] = [self.stage_index + 1] if not self.is_last else [] def _create_grad_recv_info( self, act_send_info: dict, ) -> tuple[_RecvInfo, ...]: - grad_recv_info: tuple[_RecvInfo, ...] = () + grad_recv_infos: list[_RecvInfo] = [] if not self.is_last: + # Ensure output_grads metadata is available + if self._stage_meta.output_grads is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: output_grads metadata is required for " + f"creating grad recv info. Ensure backward metadata is populated." + ) + # Receiving gradients from multiple sources is not supported # hence we only take the first destination - grad_recv_info = tuple( - _RecvInfo( - f"recv_grad_for_{self.stage_index}_from_{dst_list[0]}", - dst_list[0], - _make_tensor_from_meta(self.get_outputs_meta()[idx], self.device), + # Use a helper function to safely extract the metadata + output_grads = self._stage_meta.output_grads + for idx, dst_list in act_send_info.items(): + if dst_list is None: + raise PipeliningMetadataError( + f"Stage {self.stage_index}: output {idx} is not sent to any stage." + ) + src = dst_list[0] + grad_meta = output_grads[idx] + grad_recv_infos.append( + _RecvInfo( + input_name=f"recv_grad_for_{self.stage_index}_from_{src}", + source=src, + buffer=_make_tensor_from_meta(grad_meta, self.device) + if grad_meta + else None, + tensor_meta=grad_meta, + ) ) - for idx, dst_list in act_send_info.items() - ) - return grad_recv_info + return tuple(grad_recv_infos) diff --git a/torch/distributed/rpc/__init__.py b/torch/distributed/rpc/__init__.py index adf901d6b6e3e..d9cd701e41232 100644 --- a/torch/distributed/rpc/__init__.py +++ b/torch/distributed/rpc/__init__.py @@ -87,6 +87,7 @@ def is_available() -> bool: rendezvous_iterator: Generator[tuple[Store, int, int], None, None] __all__ += ["init_rpc", "BackendType", "TensorPipeRpcBackendOptions"] + # pyrefly: ignore [unresolvable-dunder-all] __all__ = __all__ + api.__all__ + backend_registry.__all__ # noqa: PLE0605 def init_rpc( diff --git a/torch/distributed/run.py b/torch/distributed/run.py index 16ef02495e205..87d00472edf4e 100644 --- a/torch/distributed/run.py +++ b/torch/distributed/run.py @@ -747,7 +747,7 @@ def determine_local_world_size(nproc_per_node: str): return int(nproc_per_node) except ValueError as e: if nproc_per_node == "cpu": - num_proc = os.cpu_count() + num_proc = torch._utils.cpu_count() device_type = "cpu" elif nproc_per_node == "gpu": if not torch.cuda.is_available(): @@ -769,7 +769,7 @@ def determine_local_world_size(nproc_per_node: str): num_proc = torch.accelerator.device_count() device_type = torch.accelerator.current_accelerator().type # type: ignore[union-attr] else: - num_proc = os.cpu_count() + num_proc = torch._utils.cpu_count() device_type = "cpu" else: raise ValueError( @@ -828,7 +828,7 @@ def _get_logs_specs_class(logs_specs_name: str | None) -> type[LogsSpecs]: ) logger.info( - "Using logs_spec '%s' mapped to %s", logs_specs_name, str(logs_specs_cls) + "Using logs_spec '%s' mapped to %s", logs_specs_name, logs_specs_cls ) else: logs_specs_cls = DefaultLogsSpecs diff --git a/torch/distributed/tensor/__init__.py b/torch/distributed/tensor/__init__.py index 8e1be826be58f..8795890f577fb 100644 --- a/torch/distributed/tensor/__init__.py +++ b/torch/distributed/tensor/__init__.py @@ -15,6 +15,7 @@ zeros, ) from torch.distributed.tensor.placement_types import ( + _StridedShard, Partial, Placement, Replicate, @@ -63,6 +64,7 @@ Partial, Replicate, Shard, + _StridedShard, ] ) diff --git a/torch/distributed/tensor/_api.py b/torch/distributed/tensor/_api.py index b059af927271b..00b984422926d 100644 --- a/torch/distributed/tensor/_api.py +++ b/torch/distributed/tensor/_api.py @@ -150,13 +150,12 @@ def backward(ctx, grad_output: torch.Tensor | None): # type: ignore[override] ), ) return ( - # pyrefly: ignore [bad-argument-type] - DTensor( - # pyrefly: ignore [bad-argument-count] + DTensor.from_local( grad_output, - grad_spec, - # pyrefly: ignore [unexpected-keyword] - requires_grad=grad_output.requires_grad, + grad_spec.device_mesh, + grad_spec.placements, + shape=grad_spec.shape, + stride=grad_spec.stride, ), None, ) @@ -361,7 +360,12 @@ def __tensor_flatten__(self): protocol to inform how to flatten a DTensor to local tensor for PT2 tracing """ - return ["_local_tensor"], (self._spec, self.requires_grad) + return ["_local_tensor", "device_mesh"], ( + self._spec.placements, + self._spec.tensor_meta, + self._spec.shard_order, + self.requires_grad, + ) @staticmethod def __tensor_unflatten__(inner_tensors, flatten_spec, outer_size, outer_stride): @@ -370,16 +374,18 @@ def __tensor_unflatten__(inner_tensors, flatten_spec, outer_size, outer_stride): "Expecting spec to be not None from `__tensor_flatten__` return value!" ) local_tensor = inner_tensors["_local_tensor"] - spec, requires_grad = flatten_spec + mesh = inner_tensors["device_mesh"] + placements, old_tensor_meta, shard_order, requires_grad = flatten_spec unflatten_tensor_meta = TensorMeta( shape=outer_size, stride=outer_stride, - dtype=spec.tensor_meta.dtype, + dtype=old_tensor_meta.dtype, ) unflatten_spec = DTensorSpec( - spec.mesh, - spec.placements, + mesh, + placements, tensor_meta=unflatten_tensor_meta, + shard_order=shard_order, ) # pyrefly: ignore [bad-argument-type] return DTensor( @@ -411,10 +417,10 @@ def __coerce_same_metadata_as_tangent__(self, flatten_spec, expected_type=None): if expected_type is not None: return None - (spec, _) = flatten_spec # Result of tensor_flatten() + (placements, _, _, _) = flatten_spec return self.redistribute( device_mesh=self.device_mesh, - placements=spec.placements, + placements=placements, ) @classmethod @@ -654,8 +660,9 @@ def redistribute( Returns: A :class:`DTensor` object - .. note:: ``redistribute`` is differentiable, which means user do not need to worry about - the backward formula of the redistribute operation. + .. note:: ``redistribute`` is twice-differentiable, which means user do not need to worry about + the backward formula of the redistribute operation, or its compatibility with autograd for + second-order gradients. Higher-order differentiation has not been tested (but may work). .. note:: ``redistribute`` currently only supports redistributing DTensor on the same DeviceMesh, Please file an issue if you need to redistribute DTensor to different DeviceMesh. @@ -799,19 +806,6 @@ def __get_tensor_shard__(self, index): else: raise RuntimeError("Unsupported tensor type!") - @classmethod - def __metadata_guard__( - cls, orig: tuple[DTensorSpec, bool], other: tuple[DTensorSpec, bool] - ) -> bool: - # TODO - delete this - This is now unused after the PR - - # https://github.com/pytorch/pytorch/pull/165824 - orig_spec, orig_requires_grad = orig - other_spec, other_requires_grad = other - return ( - orig_spec._check_equals(other_spec, skip_shapes=True) - and orig_requires_grad == other_requires_grad - ) - def distribute_tensor( tensor: torch.Tensor, diff --git a/torch/distributed/tensor/_collective_utils.py b/torch/distributed/tensor/_collective_utils.py index cd2e69309c45f..08f8a8eaba3c8 100644 --- a/torch/distributed/tensor/_collective_utils.py +++ b/torch/distributed/tensor/_collective_utils.py @@ -31,6 +31,35 @@ logger = logging.getLogger(__name__) +# Opaque types must be registered before defining schemas that reference them, +# so the schema parser recognizes the type names and uses PyObjectType (which +# wraps Python objects as ConcretePyObjectHolder) instead of AnyType (which +# calls toTypeInferredIValue and fails for Python-only opaque types). +from torch.distributed.device_mesh import _register_distributed_opaque_types + + +_register_distributed_opaque_types() + +_dtensor_lib = torch.library.Library("_dtensor", "FRAGMENT") +_dtensor_lib.define( + "mesh_get_process_group(" + "torch.distributed.device_mesh.DeviceMesh mesh, int dim" + ") -> torch.distributed.distributed_c10d.ProcessGroup" +) + + +@torch.library.impl("_dtensor::mesh_get_process_group", "CompositeExplicitAutograd") +def _mesh_get_process_group_impl(mesh, dim): + return mesh.get_group(dim) + + +@torch.library.register_fake("_dtensor::mesh_get_process_group") +def _mesh_get_process_group_fake(mesh, dim): + from torch._library.fake_class_registry import maybe_unwrap_fake_script_object + + real_mesh = maybe_unwrap_fake_script_object(mesh) + return real_mesh.get_group(dim) + @torch.library.register_fake("_dtensor::shard_dim_alltoall") def _shard_dim_alltoall_meta( @@ -363,6 +392,8 @@ def _compute_placement_transition_cost( num_devices_on_mesh_dim = mesh_topo.mesh_dim_devices[mesh_dim] + # NOTE: is_shard() does not match _StridedShard; see _is_shard_like(). + # Safe today: redistribute_cost bails with inf when shard_order is None. if current_placement.is_shard() and target_placement.is_replicate(): # allgather gives larger comm bytes comm_bytes_gb *= num_devices_on_mesh_dim diff --git a/torch/distributed/tensor/_decompositions.py b/torch/distributed/tensor/_decompositions.py index b1f28c82a7359..fdf298fb9acc5 100644 --- a/torch/distributed/tensor/_decompositions.py +++ b/torch/distributed/tensor/_decompositions.py @@ -326,5 +326,5 @@ def _get_candidate_placements( options |= {Shard(i) for i in range(spec.ndim)} candidates.append(list(options)) - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] return list(itertools.product(*candidates)) diff --git a/torch/distributed/tensor/_dispatch.py b/torch/distributed/tensor/_dispatch.py index d6bf1ad188d40..d346c8e17e117 100644 --- a/torch/distributed/tensor/_dispatch.py +++ b/torch/distributed/tensor/_dispatch.py @@ -173,6 +173,11 @@ def __init__(self) -> None: aten.bernoulli.default, aten.bernoulli_.float, } + self._squeeze_inplace_ops = { + aten.squeeze_.dim, + aten.squeeze_.default, + aten.squeeze_.dims, + } self._custom_op_handlers = { aten.is_same_size.default: is_same_size_handler, aten.is_pinned.default: is_pinned_handler, @@ -390,7 +395,17 @@ def _dispatch_get_local_results_slow_path( local_results = op_call(*local_tensor_args, **op_info.local_kwargs) else: # normal case, run local sharded op computation - local_results = op_call(*local_tensor_args, **op_info.local_kwargs) + if ( + output_sharding.needs_redistribute + and output_sharding.redistribute_schema is not None + and output_sharding.redistribute_schema.op != op_call + ): + # Op was rewritten (e.g., squeeze.default → squeeze.dims) + local_results = output_sharding.redistribute_schema.op( + *local_tensor_args, **op_info.local_kwargs + ) + else: + local_results = op_call(*local_tensor_args, **op_info.local_kwargs) else: # For a non-participating device (happens on rank that does not belong to @@ -481,11 +496,9 @@ def _dispatch_fast_path_python_tail( if not isinstance(args[0], dtensor.DTensor): raise AssertionError - # NOTE: aten.squeeze_.dim is an inplace op but it also may change - # the inplace argument's tensor meta. Here we choose to special case - # this op because as far as I know this is the only inplace op that - # has such as behavior. We can extend this special case if necessary. - if op_call == aten.squeeze_.dim: + # NOTE: squeeze_ inplace ops may change the tensor's metadata + # (shape/strides). We special-case them to update the spec. + if op_call in self._squeeze_inplace_ops: # update the spec to handle tensor meta changes args[0]._spec = output_spec # use return_and_correct_aliasing to match the outer and the inner @@ -593,6 +606,13 @@ def redistribute_local_args( else: new_local_args.append(arg_spec) + # Append extra non-tensor args from rewritten schema (e.g., dims tuple). + if use_val_from_redistribute_schema: + for i in range( + len(op_info.flat_args_schema), len(flatten_args_schema_to_reshard) + ): + new_local_args.append(flatten_args_schema_to_reshard[i]) + op_info.local_args = tuple(new_local_args) def unwrap_to_op_info( diff --git a/torch/distributed/tensor/_dtensor_spec.py b/torch/distributed/tensor/_dtensor_spec.py index f40acf5fdcc18..d2456be1c7862 100644 --- a/torch/distributed/tensor/_dtensor_spec.py +++ b/torch/distributed/tensor/_dtensor_spec.py @@ -8,6 +8,7 @@ import torch from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor.placement_types import ( + _is_shard_like, _MaskPartial, _StridedShard, Partial, @@ -19,6 +20,13 @@ from torch.utils._dtype_abbrs import dtype_abbrs +# Defined here (not in placement_types.py) because decoding split_factor into +# a shard order is a DTensorSpec concern — placement_types doesn't know about +# shard orders. +class _StridedShardNotDecodableError(ValueError): + """Raised when _StridedShard split_factor cannot be decoded into a shard order.""" + + class ShardOrderEntry(NamedTuple): """ Represents how a single tensor dimension is sharded across mesh dimensions. @@ -130,10 +138,9 @@ def _normalize_placements_into_shard_order( placements, mesh ) if shard_order is None: - raise ValueError( - "use_strided_shard_as_shard_order is True, but placements: " - f"{placements} is unable to be interpreted into a corresponding " - "shard_order" + raise _StridedShardNotDecodableError( + f"_StridedShard placements {placements} cannot be decoded " + "into a corresponding shard_order" ) normalized_placements = tuple( [ @@ -163,7 +170,7 @@ def compute_default_shard_order( tensor_dim_to_mesh_dims: defaultdict[int, list[int]] = defaultdict(list) mesh_ndim = len(placements) for mesh_dim in range(mesh_ndim): - if isinstance(placements[mesh_dim], Shard | _StridedShard): + if _is_shard_like(placements[mesh_dim]): placement = placements[mesh_dim] shard_dim = placement.dim # pyrefly: ignore [missing-attribute] if shard_dim < 0: @@ -314,9 +321,7 @@ def _maybe_convert_StridedShard_to_shard_order( """ if not any(isinstance(p, _StridedShard) for p in placements): return DTensorSpec.compute_default_shard_order(placements) - max_tensor_dim = ( - max([i.dim for i in placements if isinstance(i, Shard | _StridedShard)]) + 1 - ) + max_tensor_dim = max([i.dim for i in placements if _is_shard_like(i)]) + 1 shard_order = [] tensor_dim_to_mesh_dims_order: list[list[int]] = [ @@ -324,8 +329,7 @@ def _maybe_convert_StridedShard_to_shard_order( ] for mesh_dim in reversed(range(len(placements))): cur_placement = placements[mesh_dim] - # _StridedShard may not be a subclass of Shard in the future, so write in this way: - if isinstance(cur_placement, Shard | _StridedShard): + if _is_shard_like(cur_placement): tensor_dim = cur_placement.dim mesh_dims_order = tensor_dim_to_mesh_dims_order[tensor_dim] cur_sf = 1 @@ -363,6 +367,8 @@ def _verify_shard_order(self, shard_order: ShardOrder) -> None: """Verify that the shard_order is valid and matches the placements.""" total_shard = 0 if any(isinstance(p, _StridedShard) for p in self.placements): + # _StridedShard shard_order validation not yet supported; + # the Shard-only checks below (line 390, 394) would fail. return prev_tensor_dim = -1 for entry in shard_order: @@ -558,7 +564,7 @@ def format_shard_order_str( # native dtensor-style sharding representation: map from mesh # dim to tensor dim for mesh_dim, placement in enumerate(placements): - if isinstance(placement, (Shard, _StridedShard)): + if _is_shard_like(placement): if shard_order is not None: for entry in shard_order: tensor_dim = entry.tensor_dim @@ -602,7 +608,7 @@ def ndim(self) -> int: def num_shards(self) -> int: num_shards = 1 for i, placement in enumerate(self.placements): - if placement.is_shard(): + if _is_shard_like(placement): num_shards *= self.mesh.size(i) return num_shards @@ -638,8 +644,8 @@ def dim_map(self) -> list[int]: # and int >=0 represent shard on that device mesh dim r = [-1] * self.ndim for i, placement in enumerate(self.placements): - if placement.is_shard(): - shard_dim = cast(Shard, placement).dim + if _is_shard_like(placement): + shard_dim = placement.dim if r[shard_dim] > -1: raise ValueError( f"Tensor dim {shard_dim} is already sharded on mesh dim {r[shard_dim]}," @@ -666,9 +672,8 @@ def num_shards_map(self) -> list[int]: """ r = [1] * self.ndim for i, placement in enumerate(self.placements): - if placement.is_shard(): - shard_dim = cast(Shard, placement).dim - r[shard_dim] *= self.mesh.size(i) + if _is_shard_like(placement): + r[placement.dim] *= self.mesh.size(i) return r @@ -717,7 +722,7 @@ def from_dim_map( for i, m in enumerate(dim_map): if m >= 0: placement = placements[m] - if placement.is_shard(): + if placement.is_shard(): # dim_map only produces Shard placements placement = cast(Shard, placement) raise RuntimeError( f"DeviceMesh dimension can't be mapped to two dimension of the same tensor: {i} and {placement.dim}" @@ -738,9 +743,9 @@ def is_replicated(self) -> bool: def is_sharded(self) -> bool: """ - return True if the current DTensorSpec uses Shard() placement on any mesh dims (devices) + return True if the current DTensorSpec uses Shard() or _StridedShard() placement on any mesh dims (devices) """ - return any(placement.is_shard() for placement in self.placements) + return any(_is_shard_like(placement) for placement in self.placements) def shallow_copy_with_tensor_meta( self, tensor_meta: TensorMeta | None @@ -754,4 +759,5 @@ def shallow_copy_with_tensor_meta( self.mesh, self.placements, tensor_meta=tensor_meta, + use_strided_shard_as_shard_order=self.use_strided_shard_as_shard_order, ) diff --git a/torch/distributed/tensor/_ops/_conv_ops.py b/torch/distributed/tensor/_ops/_conv_ops.py index b0bce583b4ecd..7e47bd552313d 100644 --- a/torch/distributed/tensor/_ops/_conv_ops.py +++ b/torch/distributed/tensor/_ops/_conv_ops.py @@ -1,10 +1,21 @@ # Copyright (c) Meta Platforms, Inc. and affiliates # implement matrix related ops for distributed tensor +from typing import Any + import torch from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta -from torch.distributed.tensor._op_schema import OpSchema, OutputSharding +from torch.distributed.tensor._op_schema import ( + OpSchema, + OutputSharding, + RuntimeSchemaInfo, +) +from torch.distributed.tensor._ops.single_dim_strategy import ( + _ShardingPlaceholder, + register_single_dim_strategy, +) from torch.distributed.tensor._ops.utils import register_prop_rule +from torch.distributed.tensor.placement_types import Partial, Placement, Replicate aten = torch.ops.aten @@ -138,3 +149,54 @@ def convolution_backward_rules(op_schema: OpSchema) -> OutputSharding: # for a certain output Tensor. This also applies to the conv handler # in torch/distributed/tensor/_tp_conv.py return OutputSharding([grad_input_spec, grad_weight_spec, grad_bias_spec]) + + +# Single-dim strategies for autoparallel optimizer support. +# These coexist with the prop_rules above — strategies take precedence +# in the propagation path, while the prop_rules + custom handlers in +# _tp_conv.py continue to handle runtime dispatch. + + +@register_single_dim_strategy( + [aten.convolution.default], + schema_info=RuntimeSchemaInfo(2), +) +def convolution_single_dim_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder]]: + bias_meta = args_schema[2] + # [output, input, weight, (bias)] + rule: list[Placement | _ShardingPlaceholder] = [ + _ShardingPlaceholder(0), # output + _ShardingPlaceholder(0), # input + Replicate(), # weight + ] + if bias_meta is not None: + rule.append(Replicate()) # bias + return [rule] + + +@register_single_dim_strategy( + [aten.convolution_backward.default], + schema_info=RuntimeSchemaInfo(3), +) +def convolution_backward_single_dim_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder | None]]: + bias_sizes = args_schema[3] + has_bias = bias_sizes is not None + # outputs: [grad_input, grad_weight, grad_bias] + # inputs: [grad_output, input, weight] + rule: list[Placement | _ShardingPlaceholder | None] = [ + _ShardingPlaceholder(0), # grad_input + Partial("sum"), # grad_weight + Partial("sum") if has_bias else None, # grad_bias + _ShardingPlaceholder(0), # grad_output + _ShardingPlaceholder(0), # input + Replicate(), # weight + ] + return [rule] diff --git a/torch/distributed/tensor/_ops/_math_ops.py b/torch/distributed/tensor/_ops/_math_ops.py index 15495cdb76ac9..7f47c830a4cba 100644 --- a/torch/distributed/tensor/_ops/_math_ops.py +++ b/torch/distributed/tensor/_ops/_math_ops.py @@ -33,6 +33,7 @@ ) from torch.distributed.tensor._utils import normalize_to_torch_size from torch.distributed.tensor.placement_types import ( + _is_shard_like, _StridedShard, Partial, Placement, @@ -159,7 +160,7 @@ def _replicate_dims_start_at( ) -> tuple[Placement, ...]: new_placements: list[Placement] = [] for p in placements: - if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim): + if p.is_partial() or (_is_shard_like(p) and p.dim >= start_dim): new_placements.append(Replicate()) # make it replicate else: new_placements.append(p) # keep the placement @@ -167,12 +168,16 @@ def _replicate_dims_start_at( # return new_placements which align with placements but skip the skipped_dim +# Precondition: no shard-like placement on skipped_dim (callers must +# replicate it first via replicate_reduction_dims). def _skip_dim( placements: tuple[Placement, ...], skipped_dim: int ) -> tuple[Placement, ...]: new_placements: list[Placement] = [] for p in placements: - if isinstance(p, Shard) and p.dim >= skipped_dim: + if isinstance(p, _StridedShard) and p.dim >= skipped_dim: + new_placements.append(_StridedShard(p.dim - 1, split_factor=p.split_factor)) + elif isinstance(p, Shard) and p.dim >= skipped_dim: new_placements.append(Shard(p.dim - 1)) else: new_placements.append(p) @@ -188,7 +193,7 @@ def replicate_reduction_dims( for p in placements: if p.is_partial(): new_placements.append(Replicate()) - elif isinstance(p, Shard) and p.dim in reduction_dims: + elif _is_shard_like(p) and p.dim in reduction_dims: new_placements.append(Replicate()) else: new_placements.append(p) @@ -210,7 +215,7 @@ def map_placements_after_reduction( if isinstance(placement, (Replicate, Partial)): new_placements.append(placement) else: - if not isinstance(placement, Shard | _StridedShard): + if not _is_shard_like(placement): raise AssertionError( f"Expected Shard/_StridedShard, got {type(placement)}" ) @@ -221,14 +226,14 @@ def map_placements_after_reduction( # (i.e. for the case where keepdims=True), we generate partial new_placements.append(get_placement_from_reduction_op(reduction_op)) else: - if isinstance(placement, Shard): - new_placements.append(Shard(new_shard_dim)) - else: + if isinstance(placement, _StridedShard): new_placements.append( _StridedShard( new_shard_dim, split_factor=placement.split_factor ) ) + elif isinstance(placement, Shard): + new_placements.append(Shard(new_shard_dim)) return tuple(new_placements) @@ -735,12 +740,6 @@ def foreach_max_strategy(op_schema: OpSchema) -> TupleStrategy: aten.tril.default, aten.triu.default, aten._linalg_eigh.default, - aten.upsample_bicubic2d.default, - aten.upsample_bilinear2d.default, - aten.upsample_linear1d.default, - aten.upsample_nearest2d.default, - aten.upsample_trilinear3d.default, - # TODO: support the full F.interpolate set of options. ], schema_info=RuntimeSchemaInfo(1), ) @@ -1190,393 +1189,145 @@ def nll_loss_backward_strategy(op_schema: OpSchema) -> OpStrategy: return grad_in_strategy -def _common_norm_forward_strategy( - op_schema: OpSchema, - rms_norm: bool = False, -) -> OpStrategy: - """Common forward strategy logic for layer_norm and rms_norm.""" - mesh = op_schema.get_mesh_from_args() - - if not rms_norm: - # layer_norm args: input, normalized_shape, weight, bias, eps - # for None weight and bias, their corresponding objects will - # be None as well. layer_norm_strategy returns one OpStrategy - # for the triple return values (out, mean, rstd). - if not len(op_schema.args_schema) == 5: - raise AssertionError(f"Expected 5 args, got {len(op_schema.args_schema)}") - ( - input_strategy, - normalized_shape, - weight_strategy, - bias_strategy, - _, - ) = op_schema.args_schema - else: - # rms_norm args: input, normalized_shape, weight, eps - if not len(op_schema.args_schema) == 4: - raise AssertionError(f"Expected 4 args, got {len(op_schema.args_schema)}") - ( - input_strategy, - normalized_shape, - weight_strategy, - _, - ) = op_schema.args_schema - bias_strategy = None - - # the current norm implementation requires that all - # input DTensor's sharding must be in form of OpStrategy - if not isinstance(input_strategy, OpStrategy): - raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") - if not isinstance(normalized_shape, (int, Sequence, torch.Size)): - raise AssertionError( - f"Expected int, Sequence, or torch.Size, got {type(normalized_shape)}" - ) - normalized_size = normalize_to_torch_size(normalized_shape) - - input_ndim = input_strategy.ndim - axis = input_ndim - len(normalized_size) - - # we use OpStrategy because the output values (out, mean, rstd) - # should have the same placements - output_strategy = OpStrategy([]) - for idx, input_placement_strategy in enumerate(input_strategy.strategies): - op_args_target_specs = [] - redistribute_costs = [] - input_src_spec = input_placement_strategy.output_spec - - # for the input tensor, we replicate it on the inner dims if necessary - # TODO: we can avoid forcing the redistribution once we figure out - # how to decompose layer norm - input_target_spec = DTensorSpec( - mesh=mesh, - placements=_replicate_dims_start_at(input_src_spec.placements, axis), - tensor_meta=input_src_spec.tensor_meta, - ) - op_args_target_specs.append(input_target_spec) - redistribute_costs.append( - generate_redistribute_costs(input_strategy, input_target_spec) - ) - - if weight_strategy is not None: - if not isinstance(weight_strategy, OpStrategy): - raise AssertionError( - f"Expected OpStrategy, got {type(weight_strategy)}" - ) - weight_src_spec = weight_strategy.strategies[idx].output_spec - - # for the weight tensor, we replicate it on all dims if necessary - # TODO: we can avoid forcing the redistribution once we figure out - # how to decompose layer norm - weight_target_spec = DTensorSpec( - mesh=mesh, - placements=_replicate_dims_start_at(weight_src_spec.placements), - tensor_meta=weight_src_spec.tensor_meta, - ) - op_args_target_specs.append(weight_target_spec) - redistribute_costs.append( - generate_redistribute_costs(weight_strategy, weight_target_spec) - ) - - if bias_strategy is not None: - if not isinstance(bias_strategy, OpStrategy): - raise AssertionError(f"Expected OpStrategy, got {type(bias_strategy)}") - bias_src_spec = bias_strategy.strategies[idx].output_spec - - # for the bias tensor, we replicate it on all dims if necessary - # TODO: we can avoid forcing the redistribution once we figure out - # how to decompose layer norm - bias_target_spec = DTensorSpec( - mesh=mesh, - placements=_replicate_dims_start_at(bias_src_spec.placements), - tensor_meta=bias_src_spec.tensor_meta, - ) - op_args_target_specs.append(bias_target_spec) - redistribute_costs.append( - generate_redistribute_costs(bias_strategy, bias_target_spec) - ) - - # Build per-output specs with correct tensor_meta. - # out: same shape as input, contiguous strides - # mean/rstd: shape = input_shape[:axis], contiguous strides - input_tm = input_src_spec.tensor_meta - if input_tm is None: - raise AssertionError("input_src_spec.tensor_meta is None") - input_shape = input_tm.shape - out_placements = input_target_spec.placements - - out_strides = torch._prims_common.make_contiguous_strides_for(input_shape) - out_spec = DTensorSpec( - mesh=mesh, - placements=out_placements, - tensor_meta=TensorMeta( - shape=input_shape, - stride=out_strides, - dtype=input_tm.dtype, - ), - ) - - stat_shape = torch.Size(input_shape[:axis]) - stat_strides = torch._prims_common.make_contiguous_strides_for(stat_shape) - stat_spec = DTensorSpec( - mesh=mesh, - placements=out_placements, - tensor_meta=TensorMeta( - shape=stat_shape, - stride=stat_strides, - dtype=input_tm.dtype, - ), - ) - - if rms_norm: - output_specs = (out_spec, stat_spec) - else: - output_specs = (out_spec, stat_spec, stat_spec) - - output_strategy.strategies.append( - OpSpec( - output_specs=output_specs, - input_specs=op_args_target_specs, - redistribute_cost=redistribute_costs, - ) - ) - - return output_strategy - - -@register_op_strategy( +@register_single_dim_strategy( [aten.native_layer_norm.default], schema_info=RuntimeSchemaInfo(1), ) -def layer_norm_strategy(op_schema: OpSchema) -> OpStrategy: - return _common_norm_forward_strategy(op_schema) +def layer_norm_single_dim_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder]]: + input_meta = args_schema[0] + normalized_shape = args_schema[1] + weight_meta = args_schema[2] + bias_meta = args_schema[3] + axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape)) -@register_op_strategy( + strategies: list[list[Placement | _ShardingPlaceholder]] = [] + for dim in range(axis): + # [out, mean, rstd, input, weight?, bias?] + rule: list[Placement | _ShardingPlaceholder] = [ + _ShardingPlaceholder(dim), # out + _ShardingPlaceholder(dim), # mean + _ShardingPlaceholder(dim), # rstd + _ShardingPlaceholder(dim), # input + ] + if weight_meta is not None: + rule.append(Replicate()) + if bias_meta is not None: + rule.append(Replicate()) + strategies.append(rule) + return strategies + + +@register_single_dim_strategy( [aten._fused_rms_norm.default], schema_info=RuntimeSchemaInfo(1), ) -def fused_rms_norm_strategy(op_schema: OpSchema) -> OpStrategy: - return _common_norm_forward_strategy(op_schema, rms_norm=True) - - -def _common_norm_backward_strategy( - op_schema: OpSchema, - rms_norm: bool = False, -) -> OpStrategy: - """Common backward strategy logic for layer_norm and rms_norm.""" - # backward op does not need to validate the mesh since forward op has already done it - mesh = op_schema.get_mesh_from_args(validate=False) - - if not rms_norm: - # layer_norm args: grad_out, input, normalized_shape, mean, rstd, - # weight, bias, output_mask. For None weight and bias, their - # corresponding objects will be None as well. - if not len(op_schema.args_schema) == 8: - raise AssertionError(f"Expected 8 args, got {len(op_schema.args_schema)}") - ( - grad_out_strategy, - input_strategy, - normalized_shape, - mean_strategy, - rstd_strategy, - weight_strategy, - bias_strategy, - output_mask, - ) = op_schema.args_schema - else: - # rms_norm args: grad_out, input, normalized_shape, rstd, - if not len(op_schema.args_schema) == 6: - raise AssertionError(f"Expected 6 args, got {len(op_schema.args_schema)}") - ( - grad_out_strategy, - input_strategy, - normalized_shape, - rstd_strategy, - weight_strategy, - output_mask, - ) = op_schema.args_schema - mean_strategy = None - bias_strategy = None - - if not isinstance(grad_out_strategy, OpStrategy): - raise AssertionError(f"Expected OpStrategy, got {type(grad_out_strategy)}") - if not isinstance(input_strategy, OpStrategy): - raise AssertionError(f"Expected OpStrategy, got {type(input_strategy)}") - if not isinstance(rstd_strategy, OpStrategy): - raise AssertionError(f"Expected OpStrategy, got {type(rstd_strategy)}") - if mean_strategy is not None: - if not isinstance(mean_strategy, OpStrategy): - raise AssertionError(f"Expected OpStrategy, got {type(mean_strategy)}") - - if not isinstance(normalized_shape, (int, Sequence, torch.Size)): - raise AssertionError( - f"Expected int, Sequence, or torch.Size, got {type(normalized_shape)}" - ) - normalized_size = normalize_to_torch_size(normalized_shape) - input_ndim = input_strategy.ndim - axis = input_ndim - len(normalized_size) - outer_dims = list(range(axis)) - - if not rms_norm: - if not (isinstance(output_mask, list) and len(output_mask) == 3): - raise AssertionError( - f"Expected output_mask to be list of length 3, got {type(output_mask)} " - f"of length {len(output_mask) if isinstance(output_mask, list) else 'N/A'}" - ) - else: - if not (isinstance(output_mask, list) and len(output_mask) == 2): - raise AssertionError( - f"Expected output_mask to be list of length 2, got {type(output_mask)} " - f"of length {len(output_mask) if isinstance(output_mask, list) else 'N/A'}" - ) - - # output tuple: (d_input, d_weight[, d_bias]) - out_tuple_strategy = OpStrategy([]) - for idx, input_placement_strategy in enumerate(input_strategy.strategies): - # args for OpSpec - output_specs_list: list[DTensorSpec | None] = [] - input_specs_list: list[DTensorSpec] = [] - redistribute_costs = [] - - input_src_spec = input_placement_strategy.output_spec - # arg: grad_out - # TODO: change the strategy to the following rule. - # d_input is basically a product of element-wise mul of - # grad_out, rstd, and normalized input, among which rstd - # and normalized input (x_hat) should have the same sharding - # placements, and grad_out's sharding is determined by the - # pointwise result of x_hat and weight/bias. - # TODO: now grad_out spec follows input spec. we may need - # to change it to apply a pointwise rule over grad_out, - # input, and weight. - grad_out_target_spec = DTensorSpec( - mesh=mesh, - placements=_replicate_dims_start_at(input_src_spec.placements, axis), - tensor_meta=input_src_spec.tensor_meta, - ) - input_specs_list.append(grad_out_target_spec) - redistribute_costs.append( - generate_redistribute_costs(grad_out_strategy, grad_out_target_spec) - ) - output_specs_list.append(grad_out_target_spec if output_mask[0] else None) - - # arg: input - input_target_spec = DTensorSpec( - mesh=mesh, - placements=_replicate_dims_start_at(input_src_spec.placements, axis), - tensor_meta=input_src_spec.tensor_meta, - ) - input_specs_list.append(input_target_spec) - redistribute_costs.append( - generate_redistribute_costs(input_strategy, input_target_spec) - ) +def rms_norm_single_dim_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder]]: + input_meta = args_schema[0] + normalized_shape = args_schema[1] + weight_meta = args_schema[2] - # arg: mean - if not rms_norm: - if mean_strategy is None: - raise AssertionError("Expected mean_strategy to not be None") - mean_src_spec = mean_strategy.strategies[idx].output_spec - input_specs_list.append(mean_src_spec) - redistribute_costs.append([0.0 for _ in mean_strategy.strategies]) - - # arg: rstd - rstd_src_spec = rstd_strategy.strategies[idx].output_spec - input_specs_list.append(rstd_src_spec) - redistribute_costs.append([0.0 for _ in rstd_strategy.strategies]) - - def _add_target_input_spec(strategy) -> DTensorSpec: - # shared logic for setting the weight and bias target input specs - if not isinstance(strategy, OpStrategy): - raise AssertionError(f"Expected OpStrategy, got {type(strategy)}") - src_spec = strategy.strategies[idx].output_spec - # no need to redistribute since they should be replicated in forward pass - input_specs_list.append(src_spec) - redistribute_costs.append([0.0 for _ in strategy.strategies]) - return src_spec - - # arg: weight - # d_weight = sum(grad_out * (input - mean) / rstd, outer_dim, keepdim=False) - # For RMS norm, mean is 0, so it's just: sum(grad_out * input / rstd, outer_dim, keepdim=False) - if weight_strategy is not None: - weight_src_spec = _add_target_input_spec(weight_strategy) - # TODO: now d_weight spec follows input spec w/ a reduction. - # we may need to change to a pointwise rule over grad_out and - # input, then apply a reduction. - inp_placements = _replicate_dims_start_at(input_src_spec.placements, axis) - reduce_dims_map = _infer_reduce_dims_map( - outer_dims, input_src_spec.ndim, False - ) - out_placements = map_placements_after_reduction( - inp_placements, outer_dims, reduce_dims_map, "sum" - ) - weight_out_spec = DTensorSpec( - mesh=mesh, - placements=out_placements, - tensor_meta=weight_src_spec.tensor_meta, - ) - output_specs_list.append(weight_out_spec if output_mask[1] else None) - else: - if not rms_norm: - error_msg = "output_mask[1] should not be `True` while weight argument is `None` in native_layer_norm_backward." - else: - error_msg = "output_mask[1] should not be `True` while weight argument is `None` in _fused_rms_norm_backward." - if output_mask[1] is not False: - raise AssertionError(error_msg) - output_specs_list.append(None) - - # arg: bias - # d_bias = sum(grad_out, outer_dim, keepdim=False) - if not rms_norm: - if bias_strategy is not None: - bias_src_spec = _add_target_input_spec(bias_strategy) - # d_bias spec follows a reduction over grad_out - inp_placements = _replicate_dims_start_at( - grad_out_target_spec.placements, axis - ) - reduce_dims_map = _infer_reduce_dims_map( - outer_dims, grad_out_target_spec.ndim, False - ) - out_placements = map_placements_after_reduction( - inp_placements, outer_dims, reduce_dims_map, "sum" - ) - bias_out_spec = DTensorSpec( - mesh=mesh, - placements=out_placements, - tensor_meta=bias_src_spec.tensor_meta, - ) - output_specs_list.append(bias_out_spec if output_mask[2] else None) - else: - if output_mask[2] is not False: - raise AssertionError( - "output_mask[2] should not be `True` while bias argument is `None` in native_layer_norm_backward." - ) - output_specs_list.append(None) + axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape)) - out_tuple_strategy.strategies.append( - OpSpec( - output_specs=tuple(output_specs_list), - input_specs=input_specs_list, - redistribute_cost=redistribute_costs, - ) - ) - - return out_tuple_strategy + strategies: list[list[Placement | _ShardingPlaceholder]] = [] + for dim in range(axis): + # [out, rrms, input, weight?] + rule: list[Placement | _ShardingPlaceholder] = [ + _ShardingPlaceholder(dim), # out + _ShardingPlaceholder(dim), # rrms + _ShardingPlaceholder(dim), # input + ] + if weight_meta is not None: + rule.append(Replicate()) + strategies.append(rule) + return strategies -@register_op_strategy( +@register_single_dim_strategy( [aten.native_layer_norm_backward.default], schema_info=RuntimeSchemaInfo(2), ) -def layer_norm_bwd_strategy(op_schema: OpSchema) -> OpStrategy: - return _common_norm_backward_strategy(op_schema) +def layer_norm_bwd_single_dim_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder | None]]: + input_meta = args_schema[1] + normalized_shape = args_schema[2] + # mean = args_schema[3], rstd = args_schema[4] + weight_meta = args_schema[5] + bias_meta = args_schema[6] + + axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape)) + + strategies: list[list[Placement | _ShardingPlaceholder | None]] = [] + for dim in range(axis): + # outputs: [d_input, d_weight, d_bias] — always 3 per schema + # d_weight/d_bias use None when weight/bias are None + rule: list[Placement | _ShardingPlaceholder | None] = [ + _ShardingPlaceholder(dim), # d_input + Partial("sum") if weight_meta is not None else None, # d_weight + Partial("sum") if bias_meta is not None else None, # d_bias + ] + # inputs: [grad_out, input, mean, rstd, weight?, bias?] + rule.extend( + [ + _ShardingPlaceholder(dim), # grad_out + _ShardingPlaceholder(dim), # input + _ShardingPlaceholder(dim), # mean + _ShardingPlaceholder(dim), # rstd + ] + ) + if weight_meta is not None: + rule.append(Replicate()) + if bias_meta is not None: + rule.append(Replicate()) + strategies.append(rule) + return strategies -@register_op_strategy( + +@register_single_dim_strategy( [aten._fused_rms_norm_backward.default], schema_info=RuntimeSchemaInfo(2), ) -def fused_rms_norm_bwd_strategy(op_schema: OpSchema) -> OpStrategy: - return _common_norm_backward_strategy(op_schema, rms_norm=True) +def rms_norm_bwd_single_dim_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder | None]]: + input_meta = args_schema[1] + normalized_shape = args_schema[2] + # rstd = args_schema[3] + weight_meta = args_schema[4] + + axis = len(input_meta.shape) - len(normalize_to_torch_size(normalized_shape)) + + strategies: list[list[Placement | _ShardingPlaceholder | None]] = [] + for dim in range(axis): + # outputs: [d_input, d_weight] — always 2 per schema + # d_weight uses None when weight is None + # inputs: [grad_out, input, rstd, weight?] + rule: list[Placement | _ShardingPlaceholder | None] = [ + _ShardingPlaceholder(dim), # d_input + Partial("sum") if weight_meta is not None else None, # d_weight + _ShardingPlaceholder(dim), # grad_out + _ShardingPlaceholder(dim), # input + _ShardingPlaceholder(dim), # rstd + ] + if weight_meta is not None: + rule.append(Replicate()) + strategies.append(rule) + + return strategies def sort_strategy(op_schema: OpSchema, sort_dim: int) -> OpStrategy: @@ -1876,3 +1627,123 @@ def linalg_cross_strategy( continue strategies.append([_ShardingPlaceholder(dim)] * 3) return strategies + + +# --------------------------------------------------------------------------- +# Interpolation / upsample / pooling ops +# +# These ops operate on spatial dims and are safely shardable on batch (dim 0) +# and channel (dim 1). grid_sampler is batch-only because the grid tensor has +# no channel dimension. +# --------------------------------------------------------------------------- + + +@register_single_dim_strategy( + [ + aten.upsample_nearest1d.default, + aten.upsample_nearest2d.default, + aten.upsample_nearest3d.default, + aten._upsample_nearest_exact1d.default, + aten._upsample_nearest_exact2d.default, + aten._upsample_nearest_exact3d.default, + aten._upsample_bilinear2d_aa.default, + aten.upsample_bicubic2d.default, + aten.upsample_bilinear2d.default, + aten.upsample_linear1d.default, + aten.upsample_trilinear3d.default, + ], + schema_info=RuntimeSchemaInfo(1), +) +def interp_upsample_1out_1in_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder]]: + # 1 output + 1 input = 2 placements; shard on batch (0) and channel (1) + # Upsample is a linear transformation so Partial(sum/avg) is valid. + return [ + [_ShardingPlaceholder(0)] * 2, + [_ShardingPlaceholder(1)] * 2, + [Partial("sum"), Partial("sum")], + [Partial("avg"), Partial("avg")], + ] + + +@register_single_dim_strategy( + [ + aten.max_unpool2d.default, + aten.max_unpool3d.default, + aten._adaptive_avg_pool2d_backward.default, + ], + schema_info=RuntimeSchemaInfo(1), +) +def interp_pool_1out_2in_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder]]: + # 1 output + 2 inputs = 3 placements; shard on batch (0) and channel (1) + return [ + [_ShardingPlaceholder(0)] * 3, + [_ShardingPlaceholder(1)] * 3, + ] + + +@register_single_dim_strategy( + [aten.max_pool2d_with_indices_backward.default], + schema_info=RuntimeSchemaInfo(1), +) +def pool_backward_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder]]: + # max_pool2d_with_indices_backward(grad_output, self, ..., indices) -> grad_input + # 1 output + 3 tensor inputs = 4 placements + # Order: [output, grad_output, self, indices] + input_meta = cast(TensorMeta, args_schema[0]) + strategies: list[list[Placement | _ShardingPlaceholder]] = [ + [_ShardingPlaceholder(0)] * 4, + ] + if len(input_meta.shape) >= 4: # batched: (N, C, H, W) + strategies.append([_ShardingPlaceholder(1)] * 4) + # The backward is linear in grad_output, so P(sum/avg) pass through. + # indices must be replicated (integer positions, not reducible). + # self is only used for shape, so replicate it too. + r = Replicate() + for reduce_op in ("sum", "avg"): + p = Partial(reduce_op) + strategies.append([p, p, r, r]) + return strategies + + +@register_single_dim_strategy( + [aten.grid_sampler_2d.default, aten.grid_sampler_3d.default], + schema_info=RuntimeSchemaInfo(1), +) +def grid_sampler_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder]]: + # grid_sampler_{2,3}d(input[N,C,...], grid[N,...,{2,3}]) -> output[N,C,...] + # grid has no channel dim, so only batch sharding applies to both inputs. + # Linear in input: P(sum/avg) on input with replicated grid is valid. + return [ + [_ShardingPlaceholder(0)] * 3, + [Partial("sum"), Partial("sum"), Replicate()], + [Partial("avg"), Partial("avg"), Replicate()], + ] + + +@register_single_dim_strategy( + [aten.grid_sampler_2d_backward.default, aten.grid_sampler_3d_backward.default], + schema_info=RuntimeSchemaInfo(1), +) +def grid_sampler_backward_strategy( + op: torch._ops.OpOverload, + args_schema: tuple[Any, ...], + kwargs_schema: dict[str, Any], +) -> list[list[Placement | _ShardingPlaceholder]]: + # grid_sampler_{2,3}d_backward: 2 outputs (grad_input, grad_grid) + 3 inputs = 5 placements, batch-only + return [[_ShardingPlaceholder(0)] * 5] diff --git a/torch/distributed/tensor/_ops/_matrix_ops.py b/torch/distributed/tensor/_ops/_matrix_ops.py index b4ce70938ec94..b1bda0d15fdc2 100644 --- a/torch/distributed/tensor/_ops/_matrix_ops.py +++ b/torch/distributed/tensor/_ops/_matrix_ops.py @@ -36,6 +36,7 @@ compute_local_stride, ) from torch.distributed.tensor.placement_types import ( + _StridedShard, Partial, Placement, Replicate, @@ -62,10 +63,16 @@ def transpose_strategy(op_schema: OpSchema) -> OpStrategy: if ndim <= 1: output_placements = list(input_spec.placements) else: - output_placements = [ - Shard(1 - p.dim) if isinstance(p, Shard) else p - for p in input_spec.placements - ] + output_placements: list[Placement] = [] + for p in input_spec.placements: + if isinstance(p, _StridedShard): + output_placements.append( + _StridedShard(1 - p.dim, split_factor=p.split_factor) + ) + elif isinstance(p, Shard): + output_placements.append(Shard(1 - p.dim)) + else: + output_placements.append(p) transpose_strategy = OpSpec( output_specs=DTensorSpec( mesh=input_strategy.mesh, @@ -205,6 +212,7 @@ def _scaled_mm_scale_placement( if data_placement.dim == contracting_dim: return None return _ShardingPlaceholder(0) + # NOTE: isinstance(_, Shard) does not match _StridedShard; see _is_shard_like(). elif isinstance(data_placement, Shard): if data_placement.dim == contracting_dim: return None diff --git a/torch/distributed/tensor/_ops/_pointwise_ops.py b/torch/distributed/tensor/_ops/_pointwise_ops.py index 1e8c5b28a56e0..dfc253eff8e86 100644 --- a/torch/distributed/tensor/_ops/_pointwise_ops.py +++ b/torch/distributed/tensor/_ops/_pointwise_ops.py @@ -1,61 +1,20 @@ # Copyright (c) Meta Platforms, Inc. and affiliates -import functools -from collections.abc import Callable, Sequence -from typing import cast +from collections.abc import Callable import torch from torch._ops import OpOverload -from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta -from torch.distributed.tensor._op_schema import ( - ArgsType, - KwargsType, - OpSchema, - OpSpec, - OpStrategy, - RuntimeSchemaInfo, - StrategyType, - TupleStrategy, -) -from torch.distributed.tensor._ops._math_ops import _NormPartial +from torch.distributed.tensor._dtensor_spec import TensorMeta +from torch.distributed.tensor._op_schema import ArgsType, KwargsType, RuntimeSchemaInfo from torch.distributed.tensor._ops.single_dim_strategy import ( _ShardingPlaceholder, register_single_dim_strategy, ) -from torch.distributed.tensor._ops.utils import ( - generate_redistribute_costs, - infer_broadcast_dims_map, - map_placements_after_broadcast, - normalize_dim, - register_op_strategy, -) -from torch.distributed.tensor.placement_types import ( - _StridedShard, - Partial, - Placement, - Replicate, - Shard, -) -from torch.types import _Number -from torch.utils._typing_utils import not_none +from torch.distributed.tensor._ops.utils import infer_broadcast_dims_map +from torch.distributed.tensor.placement_types import Partial, Placement, Replicate aten = torch.ops.aten prims = torch.ops.prims -# leave the remaining pointwise_ops list here for convenience, -# Below ops are some pointwise ops that are yet to be supported, -# they might not be a complete list. -# pointwise_ops = [ -# "fake_quantize_per_channel_affine", -# "fake_quantize_per_tensor_affine", -# "floor_divide", # floor_divide is deprecated -# "frexp", # multiple output pointwise op, need to add support -# "gradient", # need investigation on this op -# "imag", # complex data type only -# "quantized_batch_norm", -# "quantized_max_pool1d", -# "quantized_max_pool2d", -# "real", # complex data type only -# ] # Linear pointwise ops, split by linearity type. unary_linear_ops = [aten.to.dtype] @@ -66,7 +25,12 @@ def _common_pointwise_single_dim_strategy( ) -> Callable[ [OpOverload, ArgsType, KwargsType], list[list[Placement | _ShardingPlaceholder]] ]: - """Factory for single-dim strategies that add partial placement rules.""" + """Factory for single-dim strategies that add partial placement rules. + + Returns strategies shaped [output, *args] only. Tensor kwarg placements + (e.g. ``out``, ``lr``) are appended by the wrapper in + ``_register_single_dim_pointwise``. + """ def strategy( op: OpOverload, @@ -116,46 +80,67 @@ def strategy( return strategy +def _is_list_op(op: OpOverload) -> bool: + """Returns True if op is a foreach, amp_foreach, or fused op.""" + name = op.name() + return name.startswith(("aten::_foreach_", "aten::_amp_foreach_", "aten::_fused_")) + + +# The state_steps arg of fused adam / adamw is a Replicate scalar tensor, which will be put on +# the compute_mesh of an op across all parameter groups, even when not all parameter groups +# are on the same device mesh. This idx will help avoid hitting exceptions or unnecessary +# redistribute during sharding propagation. +_FUSED_OP_SCALAR_IDX = 5 + +# Ops registered with extra Partial rules; populated by _register_single_dim_pointwise +# when partial_extra_rules is not None, to avoid double-registration from tag discovery. +_specially_registered_ops: set[OpOverload] = set() + + def _register_single_dim_pointwise( op: OpOverload, partial_extra_rules: list[list[Placement]] | None = None, static_argnum: int = 0, ) -> None: - strategy_fn = _common_pointwise_single_dim_strategy( + if partial_extra_rules is not None: + _specially_registered_ops.add(op) + inner_fn = _common_pointwise_single_dim_strategy( partial_extra_rules=partial_extra_rules # pyrefly: ignore[bad-argument-type] ) - # For .out ops, append output placement as the out kwarg placement. - # Strategy functions author [output, *args] without kwargs. The out tensor - # must match the output placement, so we duplicate strategy[0] (output). - # This makes strategies [output, *args, out_kwarg] so _get_num_tensor_inputs - # (which counts the out kwarg) computes num_outputs correctly. - if "out" in op._schema.overload_name: - inner_fn = strategy_fn - - def _out_wrapper( - op: OpOverload, - args: ArgsType, - kwargs: KwargsType, - _fn: Callable = inner_fn, - ) -> list[list[Placement | _ShardingPlaceholder]]: - strategies = _fn(op, args, kwargs) - n_tensor_args = sum(1 for a in args if isinstance(a, TensorMeta)) - n_outputs = sum(1 for r in op._schema.returns if "Tensor" in str(r.type)) - for s in strategies: - if len(s) != n_outputs + n_tensor_args: - raise AssertionError( - f"Strategy length {len(s)} != expected {n_outputs + n_tensor_args} " - f"({n_outputs} output(s) + {n_tensor_args} args) for {op}. " - f"out kwarg will be appended by infra." - ) - return [s + [s[0]] for s in strategies] - strategy_fn = _out_wrapper + # Wrap to append tensor kwarg placements in schema declaration order. + # out = output placement (s[0]); everything else (e.g. lr) = Replicate. + # TODO: move kwargs handling upstream if this works + def strategy_fn( + op: OpOverload, + args: ArgsType, + kwargs: KwargsType, + _fn: Callable = inner_fn, + ) -> list[list[Placement | _ShardingPlaceholder]]: + strategies = _fn(op, args, kwargs) + kw_names = [k for k, v in kwargs.items() if isinstance(v, TensorMeta)] + if not kw_names: + return strategies + return [ + s + [s[0] if name == "out" else Replicate() for name in kw_names] + for s in strategies + ] + + if _is_list_op(op): + schema_info = RuntimeSchemaInfo(needs_pytree=True) + else: + schema_info = RuntimeSchemaInfo(static_argnum, static_kwargkey=["out"]) + # Fused ops (e.g. _fused_adam_) have state_steps on a potentially different + # mesh; see the note in expand_to_full_mesh_op_strategy for details. + different_mesh_args: list[int] | None = None + if op.name().startswith("aten::_fused_"): + different_mesh_args = [_FUSED_OP_SCALAR_IDX] register_single_dim_strategy( op, - schema_info=RuntimeSchemaInfo(static_argnum, static_kwargkey=["out"]), + schema_info=schema_info, allow_uneven_sharding=True, allow_unbacked_sharding=True, + different_mesh_args=different_mesh_args, )(strategy_fn) @@ -171,6 +156,11 @@ def _out_wrapper( aten.sub.Tensor, aten.sub_.Tensor, aten.sub.out, + # foreach variants + aten._foreach_add.List, + aten._foreach_add_.List, + aten._foreach_sub.List, + aten._foreach_sub_.List, ] _BINARY_ADDITIVE_RULES: list[list[Placement]] = [ @@ -191,8 +181,26 @@ def _out_wrapper( _register_single_dim_pointwise(op, _BINARY_ADDITIVE_RULES) # mul: partials propagate through either arg. div: only through numerator. -binary_mul_ops = [aten.mul.Tensor, aten.mul_.Tensor, aten.mul.out] -binary_div_ops = [aten.div.Tensor, aten.div_.Tensor, aten.div.out] +binary_mul_ops = [ + aten.mul.Tensor, + aten.mul_.Tensor, + aten.mul.out, + # foreach variants + aten._foreach_mul.List, + aten._foreach_mul_.List, + aten._foreach_mul.Tensor, + aten._foreach_mul_.Tensor, +] +binary_div_ops = [ + aten.div.Tensor, + aten.div_.Tensor, + aten.div.out, + # foreach variants + aten._foreach_div.List, + aten._foreach_div_.List, + aten._foreach_div.Tensor, + aten._foreach_div_.Tensor, +] # _UNARY_LINEAR_RULES handles the scalar promotion case: Python's __mul__/__truediv__ # promote scalars to 0-dim tensors, so aten.mul.Scalar dispatches as aten.mul.Tensor @@ -220,12 +228,20 @@ def _out_wrapper( aten.div_.Scalar, aten.mul.Scalar, aten.mul_.Scalar, + # foreach variants + aten._foreach_div.Scalar, + aten._foreach_div_.Scalar, + aten._foreach_mul.Scalar, + aten._foreach_mul_.Scalar, + aten._foreach_div.ScalarList, + aten._foreach_div_.ScalarList, + aten._foreach_mul.ScalarList, + aten._foreach_mul_.ScalarList, ] for op in scalar_linear_ops: _register_single_dim_pointwise(op, _UNARY_LINEAR_RULES, static_argnum=1) -neg_ops = [aten.neg.default, aten.neg_.default] # Non-decreasing unary ops: f(max(a,b)) = max(f(a),f(b)). # Only ops that are non-decreasing on their ENTIRE domain belong here. @@ -293,6 +309,18 @@ def _out_wrapper( aten.nan_to_num.default, aten.nan_to_num_.default, aten.nan_to_num.out, + # hardshrink: x if |x|>lambd else 0. Non-decreasing on entire domain. + aten.hardshrink.default, + # I1(x) is monotonically non-decreasing for all real x. + aten.special_modified_bessel_i1.default, + # threshold(x, t, v): x if x > t else v. Non-decreasing for v <= t (the + # common case, including the default v=0, t=0). + aten.threshold.default, + # foreach variants + aten._foreach_exp.default, + aten._foreach_exp_.default, + aten._foreach_clamp_max_.Scalar, + aten._foreach_clamp_min_.Scalar, ] _NON_DECREASING_RULES: list[list[Placement]] = [ @@ -321,14 +349,47 @@ def _out_wrapper( for op in non_increasing_unary_ops: _register_single_dim_pointwise(op, _NON_INCREASING_RULES) +# Bessel K functions are strictly decreasing for x > 0 but undefined at x <= 0. +# Only P(min)->P(max) is safe: P(min) offsets add positive values to the +# non-holding rank, keeping all inputs positive. P(max) offsets subtract, +# which can push inputs to x <= 0 producing NaN. +_POSITIVE_DOMAIN_NON_INCREASING_RULES: list[list[Placement]] = [ + [Partial("max"), Partial("min")], +] + +for op in [ + aten.special_modified_bessel_k0.default, + aten.special_modified_bessel_k1.default, + aten.special_scaled_modified_bessel_k0.default, + aten.special_scaled_modified_bessel_k1.default, +]: + _register_single_dim_pointwise(op, _POSITIVE_DOMAIN_NON_INCREASING_RULES) + # neg is linear: -(A1 + A2) = -A1 + -A2 -neg_ops = [aten.neg.default, aten.neg_.default, aten.neg.out] +neg_ops = [ + aten.neg.default, + aten.neg_.default, + aten.neg.out, + # foreach variants + aten._foreach_neg.default, + aten._foreach_neg_.default, +] _NEG_RULES: list[list[Placement]] = _UNARY_LINEAR_RULES + _NON_INCREASING_RULES for op in neg_ops: _register_single_dim_pointwise(op, _NEG_RULES) +# xlog1py(x, y) = x * log1p(y). Linear in x with y replicated: +# (a+b)*log1p(y) = a*log1p(y) + b*log1p(y). +_XLOG1PY_RULES: list[list[Placement]] = [ + [Partial("sum"), Partial("sum"), Replicate()], + [Partial("avg"), Partial("avg"), Replicate()], +] + +for op in [aten.special_xlog1py.default, aten.special_xlog1py.other_scalar]: + _register_single_dim_pointwise(op, _XLOG1PY_RULES) + # All-partial-preserving unary ops: P(x)->P(x) for all x. # TODO: positive should be removed once CIA (Copy Is All) optimizes it away. @@ -383,6 +444,8 @@ def _out_wrapper( aten.maximum.default, aten.maximum.out, prims.fmax.default, + # foreach variants + aten._foreach_maximum_.List, ] _MONOTONE_MAX_PRESERVING_BINARY_BASE_RULES: list[list[Placement]] = [ @@ -412,302 +475,29 @@ def _out_wrapper( _register_single_dim_pointwise(op, _MONOTONE_MIN_PRESERVING_BINARY_BASE_RULES) -# The linear pointwise ops map, key is op, value is the type of linearity. -# Reconstructed from category lists for the existing registration path. -linear_pointwise_ops: dict[OpOverload, int] = { - aten.to.dtype: 0, - **dict.fromkeys(binary_additive_ops, 1), - **dict.fromkeys(binary_mul_ops, 2), - **dict.fromkeys(binary_div_ops, 2), - **dict.fromkeys(scalar_linear_ops, 0), - **dict.fromkeys(neg_ops, 0), -} - -pointwise_ops = [ - # please keep the entries below alphabetically sorted - aten.__ilshift__.Scalar, - aten.__ilshift__.Tensor, +# Ops that are pointwise for DTensor purposes but lack torch.Tag.pointwise. +# TODO(pianpwk): add torch.Tag.pointwise to these ops in native_functions.yaml +# so this list can be removed. +_extra_pointwise_ops: list[OpOverload] = [ aten.__irshift__.Scalar, aten.__irshift__.Tensor, - aten.__lshift__.Scalar, - aten.__lshift__.Tensor, - aten.__rshift__.Scalar, - aten.__rshift__.Tensor, aten._conj.default, - aten.abs.default, - aten.abs.out, aten.abs_.default, - aten.acos.default, - aten.acos.out, - aten.acos_.default, - aten.acosh.default, - aten.acosh.out, - aten.acosh_.default, - aten.add.Scalar, - aten.add_.Scalar, - aten.addcdiv.default, - aten.addcdiv.out, - aten.addcdiv_.default, - aten.addcmul.default, - aten.addcmul.out, - aten.addcmul_.default, - aten.angle.default, - aten.angle.out, - aten.asin.default, - aten.asin.out, - aten.asin_.default, - aten.atan2.default, - aten.atan2.out, - aten.atan2_.default, - aten.atanh.default, - aten.atanh.out, - aten.atanh_.default, - aten.bitwise_and.Scalar, - aten.bitwise_and.Scalar_Tensor, - aten.bitwise_and.Scalar_out, - aten.bitwise_and.Tensor, - aten.bitwise_and.Tensor_out, - aten.bitwise_and_.Scalar, - aten.bitwise_and_.Tensor, - aten.bitwise_left_shift.Scalar_Tensor, - aten.bitwise_left_shift.Tensor, - aten.bitwise_left_shift.Tensor_Scalar, - aten.bitwise_left_shift.Tensor_Scalar_out, - aten.bitwise_left_shift.Tensor_out, - aten.bitwise_left_shift_.Tensor, - aten.bitwise_left_shift_.Tensor_Scalar, - aten.bitwise_not.default, - aten.bitwise_not.out, - aten.bitwise_not_.default, - aten.bitwise_or.Scalar, - aten.bitwise_or.Scalar_Tensor, - aten.bitwise_or.Scalar_out, - aten.bitwise_or.Tensor, - aten.bitwise_or.Tensor_out, - aten.bitwise_or_.Scalar, - aten.bitwise_or_.Tensor, - aten.bitwise_right_shift.Scalar_Tensor, - aten.bitwise_right_shift.Tensor, - aten.bitwise_right_shift.Tensor_Scalar, - aten.bitwise_right_shift.Tensor_Scalar_out, - aten.bitwise_right_shift.Tensor_out, - aten.bitwise_right_shift_.Tensor, - aten.bitwise_right_shift_.Tensor_Scalar, - aten.bitwise_xor.Scalar, - aten.bitwise_xor.Scalar_Tensor, - aten.bitwise_xor.Scalar_out, - aten.bitwise_xor.Tensor, - aten.bitwise_xor.Tensor_out, - aten.bitwise_xor_.Scalar, - aten.bitwise_xor_.Tensor, - aten.clamp.default, - aten.clamp.Tensor, - aten.clamp.out, - aten.clamp_.default, - aten.clamp_.Tensor, - aten.clamp_min.default, - aten.clamp_max.default, - aten.clip.default, - aten.clip.out, - aten.clip_.default, - aten.conj_physical.default, - aten.conj_physical.out, - aten.conj_physical_.default, - aten.copysign.Scalar, - aten.copysign.Scalar_out, - aten.copysign.Tensor, - aten.copysign.out, aten.copysign_.Scalar, aten.copysign_.Tensor, - aten.cos.default, - aten.cos.out, - aten.cos_.default, - aten.cosh.default, - aten.cosh.out, - aten.cosh_.default, - aten.digamma.default, - aten.digamma.out, - aten.digamma_.default, - aten.div.Tensor_mode, - aten.div.out_mode, - aten.div_.Tensor_mode, - aten.eq.Tensor, - aten.eq.Tensor_out, - aten.eq.Scalar, - aten.eq.Scalar_out, - aten.erfinv.default, - aten.erfinv.out, - aten.erfinv_.default, + aten.exponential_.default, aten.float_power.Scalar, aten.float_power.Scalar_out, aten.float_power.Tensor_Scalar, aten.float_power.Tensor_Scalar_out, aten.float_power.Tensor_Tensor, aten.float_power.Tensor_Tensor_out, - aten.float_power_.Scalar, - aten.float_power_.Tensor, - aten.fmod.Scalar, - aten.fmod.Scalar_out, - aten.fmod.Tensor, - aten.fmod.Tensor_out, - aten.fmod_.Scalar, - aten.fmod_.Tensor, - aten.frac.default, - aten.frac.out, - aten.frac_.default, - aten.gcd.default, - aten.gcd.out, - aten.ge.Scalar, - aten.ge.Tensor, - aten.gelu.default, - aten.gt.Tensor, - aten.gt.Tensor_out, - aten.gt.Scalar, - aten.gt.Scalar_out, - aten.gt.Scalar, - aten.gt.Tensor, - aten.heaviside.default, - aten.heaviside.out, - aten.hypot.default, - aten.hypot.out, - aten.hypot_.default, - aten.i0.default, - aten.i0.out, - aten.i0_.default, - aten.igamma.default, - aten.igamma.out, - aten.igamma_.default, - aten.igammac.default, - aten.igammac.out, - aten.igammac_.default, - aten.isinf.default, - aten.isnan.default, - aten.isneginf.default, - aten.isneginf.out, - aten.isposinf.default, - aten.isposinf.out, - aten.ldexp.Tensor, - aten.ldexp.out, - aten.ldexp_.default, - aten.lt.Tensor, - aten.lt.Tensor_out, - aten.lt.Scalar, - aten.lt.Scalar_out, - aten.le.Scalar, - aten.le.Tensor, - aten.lerp.Scalar, - aten.lerp.Scalar_out, - aten.lerp.Tensor, - aten.lerp.Tensor_out, - aten.lerp_.Scalar, - aten.lerp_.Tensor, - aten.lgamma.default, - aten.lgamma.out, - aten.lgamma_.default, - aten.log.default, - aten.log.out, - aten.log10.default, - aten.log10.out, - aten.log10_.default, - aten.log1p.default, - aten.log1p.out, - aten.log1p_.default, - aten.log2.default, - aten.log2.out, - aten.log2_.default, - aten.log_.default, - aten.logical_and.default, - aten.logical_and.out, - aten.logical_and_.default, - aten.logical_not.default, - aten.logical_not.out, - aten.logical_not_.default, - aten.logical_or.default, - aten.logical_or.out, - aten.logical_or_.default, - aten.logical_xor.default, - aten.logical_xor.out, - aten.logical_xor_.default, - aten.logit.default, - aten.logit.out, - aten.logit_.default, - aten.masked_fill.Scalar, aten.masked_fill_.Scalar, - aten.mvlgamma.default, - aten.mvlgamma.out, - aten.mvlgamma_.default, - aten.native_dropout_backward.default, aten.native_dropout_backward.out, - aten.ne.Scalar, - aten.nextafter.default, - aten.nextafter.out, - aten.nextafter_.default, - aten.polygamma.default, - aten.polygamma.out, aten.polygamma_.default, - aten.pow.Scalar, - aten.pow.Scalar_out, - aten.pow.Tensor_Scalar, - aten.pow.Tensor_Scalar_out, - aten.pow.Tensor_Tensor, - aten.pow.Tensor_Tensor_out, - aten.pow_.Scalar, - aten.pow_.Tensor, - aten.reciprocal.default, - aten.reciprocal.out, - aten.reciprocal_.default, - aten.remainder.Scalar, - aten.remainder.Scalar_Tensor, - aten.remainder.Scalar_out, - aten.remainder.Tensor, - aten.remainder.Tensor_out, - aten.remainder_.Scalar, - aten.remainder_.Tensor, - aten.rsqrt.default, - aten.rsqrt.out, - aten.rsqrt_.default, - aten.rsub.Scalar, - aten.signbit.default, - aten.signbit.out, - aten.silu.default, - aten.silu.out, - aten.sin.default, - aten.sin.out, - aten.sin_.default, - aten.sinc.default, - aten.sinc.out, - aten.sinc_.default, - aten.sqrt.default, - aten.sqrt.out, - aten.sqrt_.default, - aten.square.default, - aten.square.out, - aten.square_.default, - aten.sub.Scalar, - aten.sub_.Scalar, - aten.tan.default, - aten.tan.out, - aten.tan_.default, - aten.true_divide.Tensor, - aten.where.self, + aten.rrelu_with_noise.default, aten.where.self_out, - aten.xlogy.OutScalar_Self, - aten.xlogy.OutScalar_Other, - aten.xlogy.OutTensor, - aten.xlogy.Scalar_Other, - aten.xlogy.Scalar_Self, - aten.xlogy.Tensor, aten.xlogy_.Scalar_Other, - aten.xlogy_.Tensor, - # backward point-wise ops - # please keep the entries below alphabetically sorted - aten.gelu_backward.default, - aten.sigmoid_backward.default, - aten.silu_backward.default, - aten.tanh_backward.default, - aten.threshold_backward.default, - # prims ops - # please keep the entries below alphabetically sorted prims.bessel_i0e.default, prims.bessel_i1.default, prims.bessel_i1e.default, @@ -715,379 +505,13 @@ def _out_wrapper( prims.bessel_j1.default, prims.div.default, prims.erfcx.default, - prims.gcd.default, prims.frexp.default, + prims.gcd.default, prims.ndtri.default, prims.ne.default, prims.spherical_bessel_j0.default, prims.zeta.default, -] - - -# Reconstruct the original linear_pointwise_ops dict for the existing registration path. -linear_pointwise_ops: dict[OpOverload, int] = { - **dict.fromkeys(unary_linear_ops, 0), - **dict.fromkeys(binary_additive_ops, 1), - **dict.fromkeys(binary_mul_ops, 2), - **dict.fromkeys(binary_div_ops, 2), - **dict.fromkeys(scalar_linear_ops, 0), - **dict.fromkeys(neg_ops, 0), -} - - -def pointwise_strategy( - op_schema: OpSchema, - linearity: int = -1, - preserve_partial: str | None = None, -) -> OpStrategy: - """Strategy for pointwise ops on the old registration path.""" - followed_strategy_index = -1 - max_shards = -1 - max_ndim = -1 - - if op_schema.is_inplace_op(): - # inplace op should follow the first arg strategy - followed_strategy = op_schema.args_schema[0] - followed_strategy_index = 0 - elif op_schema.is_out_variant_op(): - # out variant op should follow the out kwarg strategy - followed_strategy = op_schema.kwargs_schema["out"] - # out variant is technically a kwarg for the strategy to follow so it does not - # have an "index", we set it to a reasonably large number just to indicate it's - # not a valid index - followed_strategy_index = 100 - else: - # normal pointwise op, we choose to follow the arg with - # the max shards in case operands needs reshard - # in case of multiple operands with max shard, we take - # the one with the max number of dimensions - for idx, arg_strategy in enumerate(op_schema.args_schema): - if not isinstance(arg_strategy, OpStrategy): - continue - arg_max_shards = arg_strategy.max_num_shards() - arg_max_ndim = arg_strategy.ndim - if (arg_max_shards > max_shards) or ( - arg_max_shards == max_shards and arg_max_ndim > max_ndim - ): - followed_strategy_index = idx - max_shards = arg_max_shards - max_ndim = arg_max_ndim - followed_strategy = op_schema.args_schema[followed_strategy_index] - - if not isinstance(followed_strategy, OpStrategy): - raise AssertionError(f"no strategy to follow for {op_schema}!") - return common_pointwise_strategy( - op_schema.op, - op_schema.args_schema, - followed_strategy, - followed_strategy_index, - linearity, - preserve_partial=preserve_partial, - ) - - -def linear_pointwise_strategy(op_schema: OpSchema) -> StrategyType: - """ - Linear pointwise operators can propagate pending reductions. - For example, c = add(a, b); if a is pending sum, then c will be - pending sum as well without any communication overhead. - - Note that: - 1. Only unary and binary operations are supported, out variant - ops are not supported. - 2. There're multiple types of linearity, refer to the doc of - common_pointwise_strategy for more details. - """ - linearity_type = linear_pointwise_ops.get(op_schema.op, -1) - return pointwise_strategy(op_schema, linearity=linearity_type) - - -def single_mesh_dim_pointwise_strategy( - op: OpOverload, - args_schema: ArgsType, - kwargs_schema: KwargsType, - linearity: int = -1, -) -> list[list[Placement | _ShardingPlaceholder]]: - return single_mesh_dim_common_pointwise_strategy(args_schema, linearity) - - -def single_mesh_dim_linear_pointwise_strategy( - linearity: int = -1, -) -> Callable[ - [OpOverload, ArgsType, KwargsType], list[list[Placement | _ShardingPlaceholder]] -]: - return functools.partial(single_mesh_dim_pointwise_strategy, linearity=linearity) - - -def single_mesh_dim_common_pointwise_strategy( - args_schema: ArgsType, - linearity: int = -1, - scalar_tensor_idx: int | None = None, -) -> list[list[Placement | _ShardingPlaceholder]]: - # TODO rename - tensor_arg_strategies: list[TensorMeta] = [ - arg for arg in args_schema if isinstance(arg, TensorMeta) - ] - common_shape = torch.broadcast_shapes( - *[arg.shape for arg in args_schema if isinstance(arg, TensorMeta)] - ) - placements_list: list[list[Placement | _ShardingPlaceholder]] = [] - for i in range(len(common_shape)): - # Shard output dim i, and then shard the corresponding arguments if they have a corresponding (non broadcast) dim - shard_placements: list[Placement | _ShardingPlaceholder] = [ - _ShardingPlaceholder(i) - ] - for arg in tensor_arg_strategies: - common_dim_to_arg_dim = infer_broadcast_dims_map(common_shape, arg.shape) - if common_dim_to_arg_dim[i] >= 0: - shard_placements.append(_ShardingPlaceholder(common_dim_to_arg_dim[i])) - else: - shard_placements.append(Replicate()) - - placements_list.append(shard_placements) - - if linearity == 0: - # unary op (e.g. to_copy), and also binary ops like mul.scalar - # input, output can be partial - if len(tensor_arg_strategies) != 1: - raise AssertionError("expected single tensor input for linearity==0 op") - placements_list.append([Partial("sum"), Partial("sum")]) - # TODO: do i need to check scalar_tensor_index and assign a replicate to that one, or do i omit a placement for it - # TODO: can mul.scalar work with avg or only sum? i think only sum works. common_pointwise_strategy seems - # to support both. - # TODO: also, i'll be replacing 'Partial(sum)' here with some kind of 'PartialPlaceholder', not yet designed - placements_list.append([Partial("avg"), Partial("avg")]) - - elif linearity == 1: - # binary add ops - # (A1 + B1) + (A2 + B2) == (A1 + A2) + (B1 + B2) - if len(tensor_arg_strategies) != 2: - raise AssertionError("expected two tensor inputs for linearity==1 op") - placements_list.append([Partial("sum"), Partial("sum"), Partial("sum")]) - elif linearity == 2: - # binary mul ops (2 tensor inputs) - # (A * B1) + (A * B2) == A * (B1 + B2) - if len(tensor_arg_strategies) != 2: - raise AssertionError("expected two tensor inputs for linearity==2 op") - placements_list.append([Partial("sum"), Partial("sum"), Replicate()]) - placements_list.append([Partial("sum"), Replicate(), Partial("sum")]) - - # TODO: handle scalar_tensor_idx - return placements_list - - -def copy_strategy(op_schema: OpSchema) -> StrategyType: - """ - Strategy for copy_ that preserves any Partial placement. - - copy_ simply copies data and should preserve whatever Partial placement - the destination has, regardless of the reduce_op type (sum, avg, max, min, etc.). - """ - return pointwise_strategy(op_schema, preserve_partial="all") - - -def common_pointwise_strategy( - op, - args_schema: Sequence[object], - followed_strategy: OpStrategy, - followed_strategy_index: int, - linearity: int = -1, - scalar_tensor_idx: int | None = None, - preserve_partial: str | None = None, -) -> OpStrategy: - """ - Common strategy for pointwise operations. - - Args: - args_schema: Input arguments schema - followed_strategy: Strategy to follow for output placement - followed_strategy_index: Index of the strategy being followed - linearity: depending on the operator, we support different types of linearity - -1: the operation does not support linearity - 0: the unary operation that supports linearity, output propagates partial. - 1: the binary operation supports add linearity, where it requires every operand - to be partial, output propagates partial. - 2: the binary operation supports multiplicative linearity, where it requires - the primary operand to be partial, and the other operands to be replicate, - output propagates partial. - scalar_tensor_idx: Index of the Replicate scalar tensor for which we allow the mesh - to be different from the mesh of followed_strategy - preserve_partial: If set, Partial placements with this reduce_op will be preserved - through the operation (e.g., "max" for torch.maximum, "min" for torch.minimum). - """ - # handle broadcasting - common_shape = torch.broadcast_shapes( - *[arg.shape for arg in args_schema if isinstance(arg, OpStrategy)] - ) - pointwise_strategy = OpStrategy([]) - - for op_spec in followed_strategy.strategies: - spec_to_follow = op_spec.output_spec - - out_placements: list[Placement] = [] - for placement in spec_to_follow.placements: - if isinstance(placement, Shard | _StridedShard): - shard_dim = normalize_dim(placement.dim, len(spec_to_follow.shape)) - common_ndim = len(common_shape) - new_shard_dim = common_ndim - len(spec_to_follow.shape) + shard_dim - if isinstance(placement, _StridedShard): - out_placements.append( - _StridedShard( - new_shard_dim, split_factor=placement.split_factor - ) - ) - else: - out_placements.append(Shard(new_shard_dim)) - elif isinstance(placement, Partial): - is_scalar_arg = any(isinstance(arg, _Number) for arg in args_schema) - propagate_partial = False - - # ordering matters here since NormPartial is a subclass of Partial - if isinstance(placement, _NormPartial): - # explanation for args_schema[1] >= 0 can be found in summary - # https://github.com/pytorch/pytorch/pull/170035 - propagate_partial = ( - op in norm_partial_avoidable_redistribute_ops - and args_schema[1] >= 0 # pyre-ignore[unsupported-operation] - ) - - elif isinstance(placement, Partial): - propagate_partial = not ( - op in p_sum_scalar_redistribute_ops and is_scalar_arg - ) - - # Check if this partial type should be preserved - # preserve_partial="all" preserves any Partial type (used for copy_) - if preserve_partial == "all": - out_placements.append(placement) - elif preserve_partial is not None and placement.is_partial( - preserve_partial - ): - out_placements.append(placement) - # note that only partial-sum and partial-avg are supported for linearity - elif ( - linearity >= 0 - and (placement.is_partial("sum") or placement.is_partial("avg")) - and propagate_partial - ): - # propagate the partial placement - out_placements.append(placement) - else: - # clear the partial placement if op does not support linearity - # by default we just replicate the partial, need to see if this - # is optimal for all cases - out_placements.append(Replicate()) - else: - out_placements.append(placement) - - input_specs: list[DTensorSpec] = [] - redistribute_costs: list[list[float]] = [] - for input_idx, input_arg in enumerate(args_schema): - if isinstance(input_arg, OpStrategy): - input_arg_spec = input_arg.strategies[0].output_spec - - # sanity check that all args that follow the same strategy - # are on the same DeviceMesh - if input_arg.mesh != followed_strategy.mesh: - # For the scalar tensor arg in fused ops, do not follow followed_strategy; - # instead, let the input mesh and the Replicate placements propagate through. - if input_idx == scalar_tensor_idx: - if not all(p == Replicate() for p in input_arg_spec.placements): - raise AssertionError - input_arg_target_spec = DTensorSpec( - mesh=input_arg.mesh, - placements=input_arg_spec.placements, - tensor_meta=input_arg_spec.tensor_meta, - ) - input_specs.append(input_arg_target_spec) - redistribute_costs.append( - generate_redistribute_costs( - input_arg, input_arg_target_spec - ) - ) - continue - else: - raise ValueError( - f"Could not run pointwise computation across different mesh: " - f"Found {input_arg.mesh} and {followed_strategy.mesh}!" - ) - - # every arg follow the out_placements, but need to handle broadcasting - input_arg_dims_map = infer_broadcast_dims_map( - common_shape, input_arg_spec.shape - ) - - # Determine if this input should convert Partial to Replicate based on linearity - should_convert_partial = ( - linearity == 2 - and input_idx - != followed_strategy_index # Don't convert the "followed" strategy - ) - - # For preserve_partial ops, check if non-followed input has incompatible - # Partial type. If so, it must be redistributed to Replicate first. - if ( - preserve_partial is not None - and input_idx != followed_strategy_index - ): - for out_p, in_p in zip(out_placements, input_arg_spec.placements): - if ( - isinstance(out_p, Partial) - and isinstance(in_p, Partial) - and out_p != in_p - ): - should_convert_partial = True - break - - input_target_placements = map_placements_after_broadcast( - tuple(out_placements), - common_shape, - input_arg_dims_map, - partial_to_replicate=should_convert_partial, - ) - - input_arg_target_spec = DTensorSpec( - mesh=followed_strategy.mesh, - placements=input_target_placements, - tensor_meta=input_arg_spec.tensor_meta, - ) - input_specs.append(input_arg_target_spec) - redistribute_costs.append( - generate_redistribute_costs(input_arg, input_arg_target_spec) - ) - - pointwise_strategy.strategies.append( - OpSpec( - output_specs=DTensorSpec( - mesh=followed_strategy.mesh, - placements=tuple(out_placements), - ), - input_specs=input_specs, - redistribute_cost=redistribute_costs, - ) - ) - return pointwise_strategy - - -p_sum_scalar_redistribute_ops = { - aten.add.Tensor, - aten.add_.Tensor, - aten.sub.Tensor, - aten.sub_.Tensor, -} - -norm_partial_avoidable_redistribute_ops = { - aten.div.Scalar, - aten.div_.Scalar, - aten.mul.Scalar, - aten.mul_.Scalar, -} - -for op in pointwise_ops: - _register_single_dim_pointwise(op) - -# TODO: add all for_each ops -for_each_ops = [ + # foreach variants aten._foreach_abs.default, aten._foreach_abs_.default, aten._foreach_addcdiv_.Scalar, @@ -1097,147 +521,27 @@ def common_pointwise_strategy( aten._foreach_addcmul_.Scalar, aten._foreach_addcmul_.ScalarList, aten._foreach_addcmul_.Tensor, - aten._foreach_clamp_max_.Scalar, - aten._foreach_clamp_min_.Scalar, - aten._foreach_div_.List, - aten._foreach_div_.Scalar, - aten._foreach_div_.ScalarList, - aten._foreach_div_.Tensor, - aten._foreach_div.List, - aten._foreach_div.Scalar, - aten._foreach_div.ScalarList, - aten._foreach_div.Tensor, aten._foreach_lerp_.Scalar, - aten._foreach_maximum_.List, - aten._foreach_mul.Scalar, - aten._foreach_mul.ScalarList, - aten._foreach_mul.Tensor, - aten._foreach_mul.List, - aten._foreach_mul_.Scalar, - aten._foreach_mul_.ScalarList, - aten._foreach_mul_.Tensor, - aten._foreach_mul_.List, aten._foreach_pow.List, aten._foreach_pow.ScalarList, - aten._foreach_neg.default, - aten._foreach_neg_.default, aten._foreach_reciprocal_.default, aten._foreach_sub.Scalar, aten._foreach_sub_.Scalar, - aten._foreach_sub.List, - aten._foreach_sub_.List, aten._foreach_sub.ScalarList, aten._foreach_sub_.ScalarList, aten._foreach_sqrt.default, aten._foreach_sqrt_.default, aten._foreach_zero_.default, - aten._foreach_exp.default, - aten._foreach_exp_.default, aten._foreach_cos.default, aten._foreach_cos_.default, aten._foreach_log.default, aten._foreach_log_.default, aten._amp_foreach_non_finite_check_and_unscale_.default, -] - -for_each_linearity_ops = [ + # foreach linearity variants aten._foreach_add.Scalar, aten._foreach_add_.Scalar, aten._foreach_add_.ScalarList, - aten._foreach_add.List, - aten._foreach_add_.List, -] - - -def list_pointwise_strategy( - op_schema: OpSchema, linearity: bool = False -) -> StrategyType: - """ - Apply the pointwise strategy to the zipped arguments. For example, if we - run a foreach add of two lists l1 and l2, then we apply the pointwise - strategy on each pair (l1[i], l2[i]). If the first argument is a list but - the second (or later) one is a tensor, then we broadcast the tensor by - replicating it into a list with the length of the first argument. - - Args: - mesh (DeviceMesh): device mesh for pointwise ops - op_schema (OpSchema): schema of the operator to generate strategy for - linearity (bool): specify whether op(a) + op(b) = op(a + b) - - Returns: - OpStrategy: generated strategy - """ - - def args_tuple_strategies( - args_schema: tuple[object, ...], - ) -> list[TupleStrategy | None]: - first_arg = args_schema[0] - if not isinstance(first_arg, TupleStrategy): - raise AssertionError - strategy_len = len(first_arg.children) - tuple_strategies: list[TupleStrategy | None] = [] - for arg_idx, arg in enumerate(args_schema): - if isinstance(arg, TupleStrategy): - # every tuple strategy should have the same length - if len(arg.children) != strategy_len: - raise AssertionError - tuple_strategies.append(arg) - elif isinstance(arg, OpStrategy): - if arg_idx > 0: # implicitly broadcast - tuple_strategies.append( - TupleStrategy([arg for _ in range(strategy_len)]) - ) - else: - raise RuntimeError( - f"list op only supports tuple strategy! {op_schema}" - ) - else: - # insert None as placeholder so that the idx of arg is kept - tuple_strategies.append(None) - return tuple_strategies - - args_strategies = args_tuple_strategies(op_schema.args_schema) - follow_strategy: TupleStrategy = not_none(args_strategies[0]) - list_strategy: list[OpStrategy] = [] - - for child_idx, child_strtgy in enumerate(follow_strategy.children): - if not isinstance(child_strtgy, OpStrategy): - raise AssertionError - args_schema: list[OpStrategy | None] = [ - cast(OpStrategy, arg_strategy.children[child_idx]) if arg_strategy else None - for arg_strategy in args_strategies - ] - pointwise_strategy: OpStrategy = common_pointwise_strategy( - op_schema.op, - args_schema, - child_strtgy, - linearity, - scalar_tensor_idx=( - _FUSED_OP_SCALAR_IDX if op_schema.op in fused_ops else None - ), - ) - list_strategy.append(pointwise_strategy) - return TupleStrategy(list_strategy) - - -def list_linear_pointwise_strategy(op_schema: OpSchema) -> StrategyType: - """ - for each list op stratgy that supports linearity - """ - return list_pointwise_strategy(op_schema, linearity=True) - - -for op in for_each_ops: - register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))( - list_pointwise_strategy - ) - -for op in for_each_linearity_ops: - register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))( - list_linear_pointwise_strategy - ) - -fused_ops = [ + # fused optimizer ops aten._fused_adam_.default, aten._fused_adam.default, aten._fused_adam.tensor_lr, @@ -1249,16 +553,31 @@ def list_linear_pointwise_strategy(op_schema: OpSchema) -> StrategyType: ] -# The state_steps arg of fused adam / adamw is a Replicate scalar tensor, which will be put on -# the compute_mesh of an op across all parameter groups, even when not all parameter groups -# are on the same device mesh. This idx will help avoid hitting exceptions or unnecessary -# redistribute during sharding propagation. -_FUSED_OP_SCALAR_IDX = 5 +def _get_pointwise_ops_from_tag() -> list[OpOverload]: + """ + Auto-discover pointwise ops via torch.Tag.pointwise, from ops.aten, ops.prims. + """ + ops = [] + for ns in [torch.ops.aten, torch.ops.prims]: + for attr_name in dir(ns): + attr = getattr(ns, attr_name) + if isinstance(attr, torch._ops.OpOverloadPacket): + for overload_name in attr.overloads(): + op = getattr(attr, overload_name) + if torch.Tag.pointwise in op.tags: + ops.append(op) + return ops -for op in fused_ops: - register_op_strategy(op, schema_info=RuntimeSchemaInfo(needs_pytree=True))( - list_pointwise_strategy - ) + +pointwise_ops = [ + op + for op in _get_pointwise_ops_from_tag() + _extra_pointwise_ops + if op not in _specially_registered_ops +] + + +for op in pointwise_ops: + _register_single_dim_pointwise(op) def register_inductor_prims() -> None: diff --git a/torch/distributed/tensor/_ops/_tensor_ops.py b/torch/distributed/tensor/_ops/_tensor_ops.py index 5a00d9265c088..134dea27537e4 100644 --- a/torch/distributed/tensor/_ops/_tensor_ops.py +++ b/torch/distributed/tensor/_ops/_tensor_ops.py @@ -1,6 +1,6 @@ # mypy: allow-untyped-defs # Copyright (c) Meta Platforms, Inc. and affiliates -from collections.abc import Sequence, Sized +from collections.abc import Callable, Sequence, Sized from typing import cast import torch @@ -34,6 +34,7 @@ shift_shard_dims_after_remove, ) from torch.distributed.tensor.placement_types import ( + _is_shard_like, _MaskPartial, Partial, Placement, @@ -498,7 +499,7 @@ def slice_backward_rules(op_schema: OpSchema) -> OpStrategy: new_placements: list[Placement] = [] for placement in output_spec.placements: # Redistribute to replicate only if the dim is sharded and matches the slice dim - if isinstance(placement, Shard) and placement.dim == dim: + if _is_shard_like(placement) and placement.dim == dim: new_placements.append(Replicate()) else: new_placements.append(placement) @@ -516,7 +517,7 @@ def unshard_tensor_dim( ) -> tuple[Placement, ...]: """Disallow the given tensor dimension to be sharded.""" return tuple( - p if (not isinstance(p, Shard) or p.dim != dim) else Replicate() + p if (not _is_shard_like(p) or p.dim != dim) else Replicate() for p in placements ) @@ -525,10 +526,8 @@ def replicate_tensor_dim( placements: Sequence[Placement], dim: int ) -> tuple[Placement, ...]: """Force the given tensor dimension to be replicated.""" - # Not using p.is_shard() to avoid mypy complain about Placement not having - # attribute dim. return tuple( - Replicate() if p.is_partial() or isinstance(p, Shard) and p.dim == dim else p + Replicate() if p.is_partial() or (_is_shard_like(p) and p.dim == dim) else p for p in placements ) @@ -599,6 +598,70 @@ def gen_slice_scatter_strategy(op_schema: OpSchema) -> StrategyType: return slice_scatter_strategy +@register_single_dim_strategy( + [aten.select_scatter.default], + schema_info=RuntimeSchemaInfo(1), +) +def select_scatter_single_dim_strategy( + op: OpOverload, + args_schema: ArgsType, + kwargs_schema: KwargsType, +) -> list[list[Placement | _ShardingPlaceholder]]: + input_meta = args_schema[0] + if not isinstance(input_meta, TensorMeta): + raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}") + ndim = len(input_meta.shape) + dim = normalize_dim(cast(int, args_schema[2]), ndim) + # [output, self, src] — src has the select dim removed + strategies: list[list[Placement | _ShardingPlaceholder]] = [] + for d in range(ndim): + if d == dim: + continue + strategies.append( + [ + _ShardingPlaceholder(d), + _ShardingPlaceholder(d), + _ShardingPlaceholder(d if d < dim else d - 1), + ] + ) + return strategies + + +@register_single_dim_strategy( + [aten.diagonal_scatter.default], + schema_info=RuntimeSchemaInfo(1), +) +def diagonal_scatter_single_dim_strategy( + op: OpOverload, + args_schema: ArgsType, + kwargs_schema: KwargsType, +) -> list[list[Placement | _ShardingPlaceholder]]: + input_meta = args_schema[0] + if not isinstance(input_meta, TensorMeta): + raise AssertionError(f"Expected TensorMeta, got {type(input_meta)}") + ndim = len(input_meta.shape) + # schema: (self, src, offset=0, dim1=0, dim2=1) + dim1 = cast(int, args_schema[3]) if len(args_schema) > 3 else 0 + dim2 = cast(int, args_schema[4]) if len(args_schema) > 4 else 1 + dim1 = normalize_dim(dim1, ndim) + dim2 = normalize_dim(dim2, ndim) + min_d, max_d = min(dim1, dim2), max(dim1, dim2) + # [output, self, src] — src has dim1/dim2 removed and diagonal appended + strategies: list[list[Placement | _ShardingPlaceholder]] = [] + for d in range(ndim): + if d in (dim1, dim2): + continue + removed = (1 if d > min_d else 0) + (1 if d > max_d else 0) + strategies.append( + [ + _ShardingPlaceholder(d), + _ShardingPlaceholder(d), + _ShardingPlaceholder(d - removed), + ] + ) + return strategies + + @register_op_strategy(aten._local_scalar_dense.default) def replica_only_strategy(op_schema: OpSchema) -> StrategyType: """Only allow replication on the input/output.""" @@ -740,11 +803,13 @@ def merge_placement( # check each placement for the current arg placement # to see if we want to merge/adjust the placement to follow # the priority: Partial -> Shard -> Replicate + # _StridedShard.__eq__ compares both dim and split_factor, + # so two _StridedShard with different split_factor won't match here. if cur_placement == new_placement: return cur_placement if cur_placement.is_partial(): - if new_placement.is_shard(): + if _is_shard_like(new_placement): # follow new placement return new_placement elif new_placement.is_partial(): @@ -753,8 +818,8 @@ def merge_placement( else: # follow partial return cur_placement - elif cur_placement.is_shard(): - if new_placement.is_shard(): + elif _is_shard_like(cur_placement): + if _is_shard_like(new_placement): # cur/new placement are different sharding (i.e. different shard dim) # currently fallback to replicate all args return Replicate() @@ -1073,7 +1138,7 @@ def values_dim_to_output_dim(d: int) -> int: @register_single_dim_strategy( - [aten.index_put.default, aten._index_put_impl_.default], + [aten.index_put.default, aten.index_put_.default, aten._index_put_impl_.default], schema_info=RuntimeSchemaInfo(needs_pytree=True), ) def index_put_single_dim_strategy( @@ -1096,9 +1161,10 @@ def index_put_single_dim_strategy( serves as an indexing coordinate into self. Each coordinate selects a tensor element, or a slice (if non-indexed dims exist). - values is a tensor that's broadcastable to (*broadcasted_shape, *non_indexed_dim_sizes). - Each position in broadcasted_shape selects a slice of values, - and writes it into the corresponding slice of self. + values is a tensor broadcastable to the indexing output shape. + When indexed dims are consecutive starting at dim k, this shape is + (*self[:k], *broadcast_shape, *self[k+n_indexed:]). When indexed + dims are non-consecutive, it is (*broadcast_shape, *non_indexed_dims). Sharding rules (possibly conservative and incomplete): - Index tensors: always Replicate (every rank needs all coordinates). @@ -1107,10 +1173,6 @@ def index_put_single_dim_strategy( The exception is broadcasted value dimensions (size 1) - we require Replicate, but can shard self. - Additionally, we allow the full Partial rule on non-indexing tensors. - TODO(pianpwk): support non-contiguous indexed dims (None gaps in indices tuple, - e.g. (idx, None, idx)). Currently blocked by a single_dim_strategy infra bug: - _get_num_tensor_inputs counts None TupleStrategy children but args_strategy - drops them, causing a length mismatch in expand_to_full_mesh_op_strategy. """ self_meta = cast(TensorMeta, args[0]) indices_meta = cast(tuple[TensorMeta | None, ...], args[1]) @@ -1123,25 +1185,55 @@ def index_put_single_dim_strategy( values_ndim = len(values_meta.shape) # Explicitly compute the broadcast shape of the index tensors. - # We could probably derive it in a smarter way, but this is more explicit. index_shapes = [idx.shape for idx in indices_meta if idx is not None] broadcast_ndim = len(torch.broadcast_shapes(*index_shapes)) if index_shapes else 0 - # values shape = (*broadcast_shape, *non_indexed_dim_sizes) # Strategy format: [output, input, *indices, value] # The infra flattens the indices list and drops None entries, so only # non-None index tensors get a placement slot (all Replicate). + # + # Values dim mapping depends on whether indexed dims are contiguous: + # Contiguous (e.g., (None, idx0, idx1)): broadcast replaces indexed block in-place. + # values shape = (*non_indexed_before, *broadcast_shape, *non_indexed_after) + # Non-contiguous (e.g., (idx0, None, idx1)): broadcast goes to front. + # values shape = (*broadcast_shape, *non_indexed_dim_sizes) + indexed_dims_sorted = sorted(indexed_dims) + contiguous_indexed = len(indexed_dims_sorted) <= 1 or ( + indexed_dims_sorted[-1] - indexed_dims_sorted[0] + 1 == len(indexed_dims_sorted) + ) + strategies: list[list[Placement | _ShardingPlaceholder]] = [] - for values_dim in range(broadcast_ndim, values_ndim): - self_dim = non_indexed_dims[values_dim - broadcast_ndim] + for i, self_dim in enumerate(non_indexed_dims): + if contiguous_indexed and indexed_dims_sorted: + # Broadcast replaces the indexed block in-place. + first_indexed = indexed_dims_sorted[0] + if self_dim < first_indexed: + values_dim = self_dim + else: + values_dim = self_dim - n_indexed + broadcast_ndim + else: + # Broadcast goes to front (non-contiguous or no indexed dims). + values_dim = broadcast_ndim + i + + # values_dim is the position in the result tensor, but values may + # have fewer dims (right-aligned broadcasting). Convert to the + # actual values tensor dimension. + result_ndim = broadcast_ndim + len(non_indexed_dims) + values_tensor_dim = values_dim - (result_ndim - values_ndim) + + if values_tensor_dim < 0: + values_placement: Placement | _ShardingPlaceholder = Replicate() + elif values_meta.shape[values_tensor_dim] == 1: + values_placement = Replicate() + else: + values_placement = _ShardingPlaceholder(values_tensor_dim) + strategies.append( [ _ShardingPlaceholder(self_dim), _ShardingPlaceholder(self_dim), *([Replicate()] * n_indexed), - Replicate() - if values_meta.shape[values_dim] == 1 - else _ShardingPlaceholder(values_dim), + values_placement, ] ) @@ -1157,6 +1249,98 @@ def index_put_single_dim_strategy( return strategies +def _index_dim_strategy( + args_schema: ArgsType, + shard_row: Callable[[int], list[Placement | _ShardingPlaceholder]], + partial_rules: list[list[Placement | _ShardingPlaceholder]] | None = None, +) -> list[list[Placement | _ShardingPlaceholder]]: + """Common strategy for index ops that shard on all dims except the indexed dim. + + Args: + shard_row: given a dim d, returns the strategy row for sharding on that dim. + partial_rules: additional Partial passthrough strategies. + """ + self_meta = cast(TensorMeta, args_schema[0]) + ndim = len(self_meta.shape) + dim = normalize_dim(cast(int, args_schema[1]), ndim) + strategies: list[list[Placement | _ShardingPlaceholder]] = [] + for d in range(ndim): + if d != dim: + strategies.append(shard_row(d)) + if partial_rules: + strategies.extend(partial_rules) + return strategies + + +@register_single_dim_strategy( + [aten.index_fill.int_Scalar, aten.index_fill_.int_Scalar], + schema_info=RuntimeSchemaInfo(1), +) +def index_fill_scalar_single_dim_strategy( + op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType +) -> list[list[Placement | _ShardingPlaceholder]]: + # index_fill(self, dim, index, value) — fills self[..., index, ...] with scalar value. + # Partial rules: each rank fills with the same scalar v, then reduces. + # Only idempotent reduces work: avg(v,v,...,v)=v, max(v,v,...,v)=v, min(v,v,...,v)=v. + # sum and product fail: sum(v,v,...,v)=nv, product(v,v,...,v)=v^n. + return _index_dim_strategy( + args_schema, + lambda d: [ + _ShardingPlaceholder(d), # result + _ShardingPlaceholder(d), # self + Replicate(), # value (scalar, same on all ranks) + ], + [[Partial(op), Partial(op), Replicate()] for op in ("avg", "max", "min")], + ) + + +@register_single_dim_strategy( + [aten.index_fill.int_Tensor, aten.index_fill_.int_Tensor], + schema_info=RuntimeSchemaInfo(1), +) +def index_fill_tensor_single_dim_strategy( + op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType +) -> list[list[Placement | _ShardingPlaceholder]]: + # index_fill(self, dim, index, value) — fills self[..., index, ...] with 0-d tensor value. + # Partial rules: each rank fills with its partial value v_i, then reduces. + # All reduce ops work because reduce(v_0, ..., v_{n-1}) = V (the global value) + # regardless of op, since fill is a pure replacement (no mixing with self). + return _index_dim_strategy( + args_schema, + lambda d: [ + _ShardingPlaceholder(d), # result + _ShardingPlaceholder(d), # self + Replicate(), # index + Replicate(), # value + ], + [ + [Partial(op), Partial(op), Replicate(), Partial(op)] + for op in Partial.ALL_REDUCE_OPS + ], + ) + + +@register_single_dim_strategy( + [aten.index_reduce.default, aten.index_reduce_.default], + schema_info=RuntimeSchemaInfo(1), +) +def index_reduce_single_dim_strategy( + op: OpOverload, args_schema: ArgsType, kwargs_schema: KwargsType +) -> list[list[Placement | _ShardingPlaceholder]]: + # index_reduce(self, dim, index, source, reduce) — reduces source into self at index positions. + # No partial rules: reduce ops are "mean"/"amax"/"amin"/"prod", which don't match + # any Partial reduce op names ("avg"/"max"/"min"/"product"/"sum"). + return _index_dim_strategy( + args_schema, + lambda d: [ + _ShardingPlaceholder(d), # result + _ShardingPlaceholder(d), # self + Replicate(), # index + _ShardingPlaceholder(d), # source + ], + ) + + @register_op_strategy( [ aten.split.Tensor, diff --git a/torch/distributed/tensor/_ops/_view_ops.py b/torch/distributed/tensor/_ops/_view_ops.py index e6e1478308e42..29d2c3d9e35d8 100644 --- a/torch/distributed/tensor/_ops/_view_ops.py +++ b/torch/distributed/tensor/_ops/_view_ops.py @@ -1,8 +1,9 @@ # mypy: allow-untyped-defs # Copyright (c) Meta Platforms, Inc. and affiliates +import math from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass -from typing import cast +from typing import cast, NamedTuple import torch from torch import Tensor @@ -40,6 +41,13 @@ Shape = tuple[int, ...] +class ClaimedDim(NamedTuple): + """An (input_dim, output_dim) pair claimed by a mesh dim's _StridedShard rewrite.""" + + input_dim: int + output_dim: int + + @dataclass class DimSpec: """Specifies how an output dimension maps to an input dimension.""" @@ -57,12 +65,27 @@ class Singleton(DimSpec): """Output dimension is a singleton.""" -@dataclass +@dataclass(eq=False) class InputDim(DimSpec): """Output dimension maps directly to an input dimension.""" input_dim: int + def __eq__(self, other: object) -> bool: + """Raises TypeError for non-DimSpec comparisons to catch accidental + ``shard.dim == input_dim`` bugs where ``.input_dim`` was intended.""" + if isinstance(other, InputDim): + return self.input_dim == other.input_dim + if not isinstance(other, DimSpec): + raise TypeError( + f"Cannot compare InputDim with {type(other).__name__}. " + f"Did you mean to use .input_dim?" + ) + return NotImplemented + + def __hash__(self) -> int: + return hash((InputDim, self.input_dim)) + @dataclass class Broadcast(DimSpec): @@ -465,18 +488,22 @@ def dim_transpose(ndim: int, dim1: int, dim2: int) -> DimMap: return tuple(dimmap) -def dim_squeeze(shape: Shape, dim: int | None = None) -> DimMap: - # FIXME: this is wrong when dim=None and one of the dimensions - # equals size of the mesh. For example squeeze(DTensor(tensor(4), Shard[0])) could - # end up as squeeze(tensor(1)) if we have 4 devices; this would lead to - # removal of a dimension that is not actually a singleton. +def dim_squeeze(shape: Shape, dim: DimsType | None = None) -> DimMap: + # Operates on local shape; sharding_prop rewrites squeeze ops to squeeze.dims + # with only globally-singleton dims before this is called. from torch.fx.experimental.symbolic_shapes import guard_or_true + ndim = len(shape) + if dim is None: + target_dims = set(range(ndim)) + elif isinstance(dim, int): + target_dims = {normalize_dim(dim, ndim)} + else: + target_dims = set(normalize_dims(dim, ndim)) return tuple( InputDim(i) for i, s in enumerate(shape) - if guard_or_true(s > 1) - or (dim is not None and i != normalize_dim(dim, len(shape))) + if guard_or_true(s > 1) or i not in target_dims ) @@ -554,186 +581,667 @@ def propagate_shape_and_sharding( Sharding propagation follows mapped dimensions: - An output dimension that maps directly to an input dimension is sharded equally - - An output dimension that is a flattened set of input dimensions can only be - sharded if only the leftmost flattened dimension is sharded. + - An output dimension that is a flattened set of input dimensions can be sharded: + the first sharded dim stays as Shard, non-first sharded dims become _StridedShard - An output dimension that is a split of the input dimension can only be sharded if the leftmost split size is divisible by the mesh dimension """ - if not len(input_src_placements) == len(mesh_sizes): - raise AssertionError(f"{input_src_placements} != {mesh_sizes}") - # for each input dim, for each mesh dim, provides a list of possible shardable dimensions - mesh_ndim = len(mesh_sizes) - shardable_dims: dict[int, list[bool]] = {} - - # in case an input dimension disappears (e.g. collapsing, reduction) - # we cannot shard in that dimension (we need a replication fall-back rule) - seen_input_dims: set[int] = set() + propagator = _ViewShardingPropagator( + input_src_placements, global_input_shape, rule, mesh_sizes, strict_view + ) + input_tgt_placements, input_to_output_tensor_dims = propagator.analyze() + output_placements = propagator.rewrite_output_placements( + input_tgt_placements, input_to_output_tensor_dims + ) + return input_tgt_placements, output_placements - def collect_used_inputs(cmd: DimSpec) -> None: - if isinstance(cmd, InputDim): - seen_input_dims.add(cmd.input_dim) - for inp in cmd.inputs(): - collect_used_inputs(inp) - for cmd in rule: - collect_used_inputs(cmd) - for dim in range(len(global_input_shape)): - shardable_dims[dim] = [dim in seen_input_dims] * mesh_ndim +class _ViewShardingPropagator: + """Two-phase sharding propagator for view ops. + + Phase 1 — ``analyze()``: + Walks the DimMap rule and returns: + - ``input_tgt_placements``: input placements with unshardable dims + demoted to Replicate. + - ``input_to_output_tensor_dims``: maps each input tensor dim to its + output dim(s). Cardinality encodes the op type: 1→1 for InputDim, + N→1 for Flatten, 1→N for Split/unflatten. + + Phase 2 — ``rewrite_output_placements()``: + Consumes both Phase 1 outputs. Iterates mesh dims 0..n-1, maintaining: + - ``strided_shard_claimed_dims``: (input_dim, output_dim) pairs already assigned + to a mesh dim by _StridedShard rewriting. + - ``local_tensor_shapes``: global shape progressively divided by each + mesh dim's shard size. + For each surviving Shard/_StridedShard, looks up the output dim(s) and + produces the final output placement. + """ - def maybe_get_shard_mesh_dim_and_placement( - input_dim: InputDim, - ) -> tuple[int | None, Shard | _StridedShard | None]: - # if input_dim is sharded, return the mesh_dim and shard placement - for i, placement in enumerate(input_src_placements): + def __init__( + self, + input_src_placements: Sequence[Placement], + global_input_shape: Shape, + rule: DimMap, + mesh_sizes: Shape, + strict_view: bool, + ) -> None: + self.input_src_placements = input_src_placements + self.global_input_shape = global_input_shape + self.rule = rule + self.mesh_sizes = mesh_sizes + self.strict_view = strict_view + self.mesh_ndim = len(mesh_sizes) + + # shard_allowed[input_dim][mesh_dim]: whether input_dim can stay + # sharded on mesh_dim. Populated by _analyze_dim and its helpers. + self.shard_allowed: dict[int, list[bool]] = {} + # Mesh dims whose _StridedShard has already been matched to an output dim. + # Populated by _analyze_split. + self.matched_strided_mesh_dims: set[int] = set() + + # ------------------------------------------------------------------ + # Public API: analyze → rewrite_output_placements + # ------------------------------------------------------------------ + + def analyze( + self, + ) -> tuple[Sequence[Placement], dict[int, list[int]]]: + """Phase 1: walk the DimMap rule, return (input_tgt_placements, input_to_output_tensor_dims).""" + input_dims_in_rule = self._input_dims_in_rule(self.rule) + + # Default: shardable if the dim appears in the rule. Refined by _analyze_*. + for dim in range(len(self.global_input_shape)): + self.shard_allowed[dim] = [dim in input_dims_in_rule] * self.mesh_ndim + + # Walk the rule to refine shard_allowed and build input_to_output_tensor_dims. + # + # Flatten example: view([2, 3, 4], [6, 4]) + # rule = (Flatten(InputDim(0), InputDim(1)), InputDim(2)) + # output_dim=0 (Flatten): hits the isinstance(cmd, Flatten) branch. + # Maps input dims 0 and 1 to output dim 0. Result: {0: [0], 1: [0]} + # output_dim=1 (InputDim(2)): hits the len(in_dims) > 0 branch. + # Maps input dim 2 to output dim 1. Result: {0: [0], 1: [0], 2: [1]} + # + # Split example: view([6], [2, 3]) + # rule = (Split(InputDim(0), (2,3), 0), Split(InputDim(0), (2,3), 1)) + # output_dim=0 (split_id=0): hits the len(in_dims) > 0 branch. + # Maps input dim 0 to output dim 0. Result: {0: [0]} + # output_dim=1 (split_id=1): hits the isinstance(cmd, Split) branch + # because _analyze_split returns [] for split_id>0. Chases root + # InputDim(0) and appends output dim 1. Result: {0: [0, 1]} + input_to_output_tensor_dims: dict[int, list[int]] = {} + for output_dim, cmd in enumerate(self.rule): + in_dims = self._analyze_dim(cmd) + if isinstance(cmd, Flatten): + for in_dim in in_dims: + if in_dim.input_dim in input_to_output_tensor_dims: + raise AssertionError( + f"Input dim {in_dim.input_dim} already mapped to output dims " + f"{input_to_output_tensor_dims[in_dim.input_dim]}" + ) + input_to_output_tensor_dims[in_dim.input_dim] = [output_dim] + elif len(in_dims) > 0: + # InputDim (identity) or Split(split_id=0). + in_dim = in_dims[0] + if in_dim.input_dim not in input_to_output_tensor_dims: + input_to_output_tensor_dims[in_dim.input_dim] = [output_dim] + else: + input_to_output_tensor_dims[in_dim.input_dim].append(output_dim) + elif isinstance(cmd, Split): + # Split(split_id>0): _analyze_split returned [], so chase the + # root input dim and append this output dim to its existing entry. + # + # Flatten+Split example: view([2, 3], [3, 2]) + # rule = (Split(Flatten(InputDim(0), InputDim(1)), (3,2), 0), + # Split(Flatten(InputDim(0), InputDim(1)), (3,2), 1)) + # output_dim=0 (split_id=0): same as Split example above. + # Result: {0: [0]} + # output_dim=1 (split_id=1): same as Split example, but + # the chase unwraps the inner Flatten to find InputDim(0). + # Result: {0: [0, 1]} + root_spec = cmd.input_dim + while isinstance(root_spec, (Flatten, Split)): + if isinstance(root_spec, Flatten): + # _analyze_flatten always returns input_dims[0] as + # the first element (either as the only shardable dim + # in non-strict mode, or as the fallback when nothing + # is sharded), so split_id=0 uses it as the key in + # input_to_output_tensor_dims. Use [0] here to match. + root_spec = root_spec.input_dims[0] + else: + root_spec = root_spec.input_dim + root = root_spec if isinstance(root_spec, InputDim) else None + if root is not None and root.input_dim in input_to_output_tensor_dims: + input_to_output_tensor_dims[root.input_dim].append(output_dim) + + input_tgt_placements: list[Placement] = [] + for mesh_dim, p in enumerate(self.input_src_placements): if ( - isinstance(placement, Shard | _StridedShard) - and placement.dim == input_dim.input_dim + isinstance(p, Shard | _StridedShard) + and not self.shard_allowed[p.dim][mesh_dim] ): - return i, placement - return None, None + input_tgt_placements.append(Replicate()) + else: + input_tgt_placements.append(p) + return input_tgt_placements, input_to_output_tensor_dims + + def rewrite_output_placements( + self, + input_tgt_placements: Sequence[Placement], + input_to_output_tensor_dims: dict[int, list[int]], + ) -> list[Placement]: + """Phase 2: consume analyze() outputs, return final output placements.""" + # (input_dim, output_dim) pairs claimed by earlier mesh dims + # (via _rewrite_strided_shard), to avoid double-assignment. + strided_shard_claimed_dims: set[ClaimedDim] = set() + # Starts as global_input_shape; each mesh dim divides its sharded dim. + local_tensor_shapes: list[int] = list(self.global_input_shape) + + output_placements: list[Placement] = [] + # Process mesh dims in order; _rewrite_*_shard relies on this for + # truncating division safety in local_tensor_shapes. + for mesh_dim, p in enumerate(input_tgt_placements): + if isinstance(p, Shard): + placement, local_tensor_shapes = self._rewrite_plain_shard( + p, + mesh_dim, + input_tgt_placements, + strided_shard_claimed_dims, + local_tensor_shapes, + input_to_output_tensor_dims, + ) + output_placements.append(placement) + elif isinstance(p, _StridedShard): + placement, local_tensor_shapes = self._rewrite_strided_shard( + p, + mesh_dim, + input_tgt_placements, + strided_shard_claimed_dims, + local_tensor_shapes, + input_to_output_tensor_dims, + ) + output_placements.append(placement) + else: + output_placements.append(p) + return output_placements + + # ------------------------------------------------------------------ + # Analysis phase helpers + # ------------------------------------------------------------------ + + @staticmethod + def _input_dims_in_rule(rule: DimMap) -> set[int]: + """Walk the DimMap rule tree and return all input dim indices that appear in it.""" + seen: set[int] = set() + + def _walk(cmd: DimSpec) -> None: + if isinstance(cmd, InputDim): + seen.add(cmd.input_dim) + for inp in cmd.inputs(): + _walk(inp) + + for cmd in rule: + _walk(cmd) + return seen + + def _find_plain_shard( + self, input_dim: InputDim + ) -> tuple[int | None, Shard | _StridedShard | None]: + """Find the mesh dim with a plain Shard on ``input_dim``. - # NOTE: This function has three responsibilities: - # 1. determine "theoretically" if an output dimension can be sharded, i.e. fill the shardable_dims map - # 2. determine "theoretically" the corresponding input dimension to shard on, via return value - # 3. throw an error when strict_view is enabled and we cannot shard an output dimension - # 1 and 2 doesn't require the info of whether current input is sharded. - # 3 requires that info, to decide whether we can error out. Maybe we can refactor - # to make this function purely "theoretical". - def get_in_dim_to_shard(cmd: DimSpec) -> InputDim | None: - from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true + Only matches Shard, not _StridedShard. Used by both _analyze_flatten + and _analyze_split. _find_shard_for_split is the counterpart that + also matches _StridedShard with split_factor validation. + """ + for mesh_dim, placement in enumerate(self.input_src_placements): + if isinstance(placement, Shard) and placement.dim == input_dim.input_dim: + return mesh_dim, placement + return None, None - if isinstance(cmd, InputDim): - return cmd - elif isinstance(cmd, Flatten): - for i, dim in enumerate(cmd.input_dims): - # so far all Flatten is always composed of InputDims; revisit this if needed - if not isinstance(dim, InputDim): - raise AssertionError(f"Expected InputDim, got {type(dim)}") - can_shard_dim = True - shard_mesh_dim, shard_placement = ( - maybe_get_shard_mesh_dim_and_placement(dim) + def _find_shard_for_split( + self, + current_dim: int, + cmd: Split, + placements: Sequence[Placement], + ) -> tuple[int | None, Shard | _StridedShard | None]: + """Find the mesh dim and placement for an input dim in Split ops. + + Matches both Shard and _StridedShard: + - Shard: plain unflatten, e.g. [6] Shard(0) → [2, 3]. + - _StridedShard: unflatten after a prior flatten that produced + _StridedShard, e.g. [2,3,4] Shard(1) → flatten → [6,4] + _StridedShard(0,sf=2) → unflatten → [2,3,4]. Validates that + the split_factor matches the expected value for this split_id. + """ + for mesh_dim, placement in enumerate(placements): + if not isinstance(placement, Shard | _StridedShard): + continue + if placement.dim != current_dim: + continue + if mesh_dim in self.matched_strided_mesh_dims: + continue + + if isinstance(placement, _StridedShard): + expected_sf = self._expected_split_factor( + cmd, current_dim, mesh_dim, placements ) - input_sharded = shard_mesh_dim is not None - if i > 0: - can_shard_dim = False - if strict_view and input_sharded: - raise RuntimeError( - f"Attempted to flatten multiple dimensions, with dimension {dim.input_dim} being sharded. ", - "It cannot be performed without redistribution, which is disallowed by the current operator.", - ) - elif input_sharded: - if not (shard_placement is not None and shard_mesh_dim is not None): - raise AssertionError( - "Expected shard_placement and shard_mesh_dim to be not None" - ) - tensor_dim_size = global_input_shape[shard_placement.dim] - mesh_dim_size = mesh_sizes[shard_mesh_dim] + if expected_sf == placement.split_factor: + return mesh_dim, placement + else: + return mesh_dim, placement + return None, None + + def _analyze_flatten(self, cmd: Flatten) -> list[InputDim]: + """Fill self.shard_allowed for Flatten; return sharded input dims.""" + from torch.fx.experimental.symbolic_shapes import guard_or_true + + sharded_dims: list[InputDim] = [] + num_input_dims = len(cmd.input_dims) + for i, dim in enumerate(cmd.input_dims): + if not isinstance(dim, InputDim): + raise AssertionError(f"Expected InputDim, got {type(dim)}") + shard_mesh_dim, shard_placement = self._find_plain_shard(dim) + if shard_mesh_dim is None or shard_placement is None: + continue # default from analyze() already covers this + tensor_dim_size = self.global_input_shape[shard_placement.dim] + mesh_dim_size = self.mesh_sizes[shard_mesh_dim] + can_shard_dim = True + if self.strict_view: + is_last_input_dim = i == num_input_dims - 1 + if not is_last_input_dim and guard_or_true( + tensor_dim_size % mesh_dim_size != 0 + ): + raise RuntimeError( + f"Cannot flatten unevenly sharded tensor: " + f"dimension {dim.input_dim} (size {tensor_dim_size}) " + f"is not evenly divisible by mesh dimension " + f"{shard_mesh_dim} (size {mesh_dim_size}). " + f"Please redistribute the tensor before this operation." + ) + sharded_dims.append(dim) + else: + # TODO: non-strict (reshape) should allow can_shard_dim = True + # for non-first flatten dims, since strict_view already does. + # Currently forces redistribution because the rewrite phase + # wasn't originally implemented for this case. + if i == 0: + sharded_dims.append(dim) if guard_or_true(tensor_dim_size % mesh_dim_size != 0): can_shard_dim = False - if strict_view: - raise RuntimeError( - f"Attempted to flatten unevenly sharded dimension {i}, " - "which would require resharding the input. " - "Please explicitly redistribute the tensor instead." - ) - shardable_dims[dim.input_dim] = [can_shard_dim] * mesh_ndim - - if not isinstance(cmd.input_dims[0], InputDim): - raise AssertionError( - f"Expected InputDim, got {type(cmd.input_dims[0])}" - ) - return cmd.input_dims[0] - elif isinstance(cmd, Split): - in_dim = get_in_dim_to_shard(cmd.input_dim) - out_size = cmd.group_shape[cmd.split_id] - if cmd.split_id == 0 and in_dim is not None: - # we need to check that the input dimension is divisible - # by the size of the submesh we're sharding it on - # NOTE: it would be possible to shard the same input dimension - # on more than one mesh dimension. In that case, the dimension - # needs to be divisible by the product of mesh sizes. - # In order to keep the problem more tractable, we will not consider - # double resharding as a suggestion (e.g. [Shard(0), Shard(0) ]) - # but we will allow it if that's the input and it's compatible - - # 1. is this dimension shardable on each individual mesh dim? - shardable_dims[in_dim.input_dim] = [ - guard_or_false(out_size % mesh_dim_size == 0) - for mesh_dim_size in mesh_sizes - ] - - shard_mesh_dim, _ = maybe_get_shard_mesh_dim_and_placement(in_dim) - if strict_view and shard_mesh_dim is not None: - if not shardable_dims[in_dim.input_dim][shard_mesh_dim]: - raise RuntimeError( - f"Attempted to split the sharded dimension {in_dim.input_dim} into multiple subdimensions. ", - "It cannot be performed without redistribution, which is disallowed by the current operator.", - ) + else: + can_shard_dim = False + self.shard_allowed[dim.input_dim] = [can_shard_dim] * self.mesh_ndim + + if len(sharded_dims) > 0: + return sharded_dims + # No sharded dims: e.g. Flatten([InputDim(0), InputDim(1)]) where + # neither dim is sharded. Return the first input dim so that + # input_to_output_tensor_dims is populated for identity rewrites. + if not isinstance(cmd.input_dims[0], InputDim): + raise AssertionError(f"Expected InputDim, got {type(cmd.input_dims[0])}") + return [cmd.input_dims[0]] + + def _analyze_split(self, cmd: Split) -> list[InputDim]: + """Fill self.shard_allowed for Split; return shardable input dims.""" + from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true - # 2. here we special case things like [Shard(0), Shard(0)] - submesh_size = 1 - for size, shard in zip(mesh_sizes, input_src_placements): - if isinstance(shard, Shard | _StridedShard) and shard.dim == in_dim: - submesh_size *= size - if guard_or_true(out_size % submesh_size != 0): - raise AssertionError( - f"Resulting dimension size {out_size} is not divisible by its mesh dimension {submesh_size}." + in_dims = self._analyze_dim(cmd.input_dim) + if len(in_dims) == 0: + return [] + in_dim = in_dims[0] + out_size = cmd.group_shape[cmd.split_id] + shard_mesh_dim, input_src_placement = self._find_shard_for_split( + in_dim.input_dim, cmd, self.input_src_placements + ) + # split_id == 0 sets the base shard_allowed for this input dim. + # Later split_ids (processed in subsequent rule iterations) refine + # individual mesh_dim entries via the _StridedShard branch below. + if cmd.split_id == 0: + self.shard_allowed[in_dim.input_dim] = [ + guard_or_false(out_size % mesh_dim_size == 0) + for mesh_dim_size in self.mesh_sizes + ] + plain_mesh_dim, _ = self._find_plain_shard(in_dim) + # Non-strict silently redistributes via shard_allowed=False above; + # strict raises so the user knows to redistribute before view(). + if self.strict_view and plain_mesh_dim is not None: + if not self.shard_allowed[in_dim.input_dim][plain_mesh_dim]: + raise RuntimeError( + f"Cannot unflatten unevenly sharded tensor: " + f"output dimension {cmd.split_id} (size {out_size}) " + f"is not evenly divisible by mesh dimension " + f"{plain_mesh_dim} (size {self.mesh_sizes[plain_mesh_dim]}). " + f"Please redistribute the tensor before this operation." ) + if shard_mesh_dim is not None and isinstance( + input_src_placement, _StridedShard + ): + # The last split dim doesn't require even divisibility because + # its local size is inferred: local_last = local_flat / product + # of earlier dims, and DTensor handles uneven local sizes. + # Non-last dims must be evenly divisible because they appear as + # fixed sizes in the local reshape — uneven division would make + # the stride pattern inconsistent across devices. + # E.g. [12] → [3, 4], _StridedShard targeting dim 1 (last), + # mesh=3: 4%3≠0, but local shapes [3,2],[3,1],[3,1] are valid. + is_last_split_dim = cmd.split_id == len(cmd.group_shape) - 1 + if ( + self.strict_view + and not is_last_split_dim + and guard_or_true(out_size % self.mesh_sizes[shard_mesh_dim] != 0) + ): + raise RuntimeError( + f"Cannot unflatten unevenly sharded tensor: " + f"output dimension {cmd.split_id} (size {out_size}) " + f"is not evenly divisible by mesh dimension {shard_mesh_dim} " + f"(size {self.mesh_sizes[shard_mesh_dim]}). " + f"Please redistribute the tensor before this operation." + ) + # Prevents _find_shard_for_split from matching this mesh dim + # again for a later split_id of the same Split group. + self.matched_strided_mesh_dims.add(shard_mesh_dim) + if in_dim.input_dim in self.shard_allowed: + self.shard_allowed[in_dim.input_dim][shard_mesh_dim] = ( + guard_or_false(out_size % self.mesh_sizes[shard_mesh_dim] == 0) + or is_last_split_dim + ) + # Only split_id==0 returns the input dim for input_to_output_tensor_dims. + # Later split_ids refine shard_allowed above but return [] — their + # output dims are linked via the root-input-dim chase in analyze(). + return [in_dim] if cmd.split_id == 0 else [] - # we will only shard our first component of the split - return in_dim if cmd.split_id == 0 else None + def _analyze_dim(self, cmd: DimSpec) -> list[InputDim]: + """Dispatch one DimSpec: update self.shard_allowed, return input dim(s) to shard on.""" + if isinstance(cmd, InputDim): + return [cmd] + elif isinstance(cmd, Flatten): + return self._analyze_flatten(cmd) + elif isinstance(cmd, Split): + return self._analyze_split(cmd) elif isinstance(cmd, Repeat): - in_dim = get_in_dim_to_shard(cmd.input_dim) - if in_dim is not None: - shardable_dims[in_dim.input_dim] = [False] * mesh_ndim - return None + in_dims = self._analyze_dim(cmd.input_dim) + for d in in_dims: + self.shard_allowed[d.input_dim] = [False] * self.mesh_ndim + return [] else: - return None - - # for each output dim, find the corresponding input dim in terms of sharding prop - shard_dim_map = {} - for dim, cmd in enumerate(rule): - in_dim = get_in_dim_to_shard(cmd) - if in_dim is not None: - shard_dim_map[in_dim.input_dim] = dim - - input_tgt_placements = [ - ( - Replicate() - if isinstance(p, Shard | _StridedShard) - and not shardable_dims[p.dim][mesh_dim] - else p + return [] + + # ------------------------------------------------------------------ + # Rewrite phase helpers + # ------------------------------------------------------------------ + + @staticmethod + def _is_last_shard_in_flatten_range( + mesh_dim: int, + placements: Sequence[Placement], + flatten_start: int, + flatten_end: int, + ) -> bool: + """Check if no later mesh dim shards a dim within the flatten range at or above this one. + + Uneven sharding on dim d breaks stride computation for all earlier dims + that flatten together with d. Only dims within [flatten_start, flatten_end) + matter; shards on dims outside the flatten range are independent. + + Requires: placements[mesh_dim] must be Shard or _StridedShard. + """ + p = placements[mesh_dim] + if not isinstance(p, (Shard, _StridedShard)): + raise AssertionError( + f"Expected Shard or _StridedShard at mesh_dim {mesh_dim}, got {type(p)}" + ) + tensor_dim = p.dim + return not any( + isinstance(other_p, (Shard, _StridedShard)) + and flatten_start <= other_p.dim < flatten_end + and other_p.dim >= tensor_dim + for other_p in placements[mesh_dim + 1 :] ) - for mesh_dim, p in enumerate(input_src_placements) - ] - def _rewrite_shard_dim(p: Shard | _StridedShard): + def _expected_split_factor( + self, + cmd: Split, + sharded_dim: int, + mesh_dim: int, + placements: Sequence[Placement], + ) -> int | None: + """Compute the residual split factor for ``cmd`` after earlier mesh dims. + + Starts from ``math.prod(cmd.group_shape[:cmd.split_id])`` and divides + out each earlier mesh dim that shards the same input dim. Returns + ``None`` if any earlier mesh size doesn't divide evenly. """ - Rewrite the shard dim to the corresponding tensor dim in output. - For ``_StridedShard``, we can safely keep the placement type and - ``split_factor`` unchanged and only rewrite the ``dim`` because: - 1. ``_StridedShard`` has no impact on sharding (i.e. how - tensor is partitioned) compared to ``Shard``. It only changes - how shards permute across the devices. - 2. ``view()`` op on DTensor strictly forbids shard redistribution - which means if ``view()`` may cause shard permutation across - devices, it should be rejected. This is enforced in today's - sharding prop for ``view()``. - 3. Since DTensor ``view()`` won't introduce any redistribution, - it's certain that ``placements`` won't change except the - inner ``dim`` attribute of ``Shard`` or ``_StridedShard``. + sf = math.prod(cmd.group_shape[: cmd.split_id]) + for m in range(mesh_dim): + other_p = placements[m] + if ( + isinstance(other_p, (_StridedShard, Shard)) + and other_p.dim == sharded_dim + ): + if sf % self.mesh_sizes[m] != 0: + return None + sf //= self.mesh_sizes[m] + return sf + + def _find_keep_ss_dim( + self, + tgt_shard_dims: list[int], + p: _StridedShard, + mesh_dim: int, + ) -> int | None: + """Find an output dim where SS stays as SS. + + Returns the first output dim whose Split can accommodate the combined + sharding (mesh_size * split_factor), or ``None`` if no dim fits. """ - if isinstance(p, _StridedShard): - return _StridedShard(shard_dim_map[p.dim], split_factor=p.split_factor) + total_shard = self.mesh_sizes[mesh_dim] * p.split_factor + if self.global_input_shape[p.dim] % total_shard != 0: + return None + shard_size = self.global_input_shape[p.dim] // total_shard + for candidate_dim in tgt_shard_dims: + cmd = self.rule[candidate_dim] + if isinstance(cmd, Split): + inner_size = math.prod(cmd.group_shape[cmd.split_id + 1 :]) + # When a Split wraps a Flatten, the per-shard chunk covers + # the sharded dim plus trailing dims flattened together. + trailing_size = 1 + if isinstance(cmd.input_dim, Flatten): + found = False + for flat_dim in cmd.input_dim.input_dims: + if not isinstance(flat_dim, InputDim): + raise AssertionError( + f"Expected InputDim, got {type(flat_dim)}" + ) + if flat_dim.input_dim == p.dim: + found = True + elif found: + trailing_size *= self.global_input_shape[flat_dim.input_dim] + flattened_shard_size = shard_size * trailing_size + if ( + flattened_shard_size >= inner_size + and flattened_shard_size % inner_size == 0 + ): + return candidate_dim + return None + + def _rewrite_plain_shard( + self, + p: Shard, + mesh_dim: int, + placements: Sequence[Placement], + strided_shard_claimed_dims: set[ClaimedDim], + local_tensor_shapes: list[int], + input_to_output_tensor_dims: dict[int, list[int]], + ) -> tuple[Placement, list[int]]: + """Given a plain Shard(dim=X) input placement on a specific mesh dim, + determine what output placement it maps to after the view op. + + For identity and unflatten, produces Shard on the output dim. For + flatten, Shard on the first flattened dim stays Shard, while Shard on + a non-first dim produces _StridedShard (consumed later by + _rewrite_strided_shard). + + Returns the output placement and a new local_tensor_shapes with this + mesh dim's division applied. + """ + # Output dims that input dim p.dim maps to, filtering out any + # already claimed by _StridedShard rewriting on earlier mesh dims. + tgt_shard_dims = [ + d + for d in input_to_output_tensor_dims[p.dim] + if ClaimedDim(p.dim, d) not in strided_shard_claimed_dims + ] + if len(tgt_shard_dims) == 0: + raise AssertionError( + f"No output dim available for Shard(dim={p.dim}) on mesh dim " + f"{mesh_dim}. All output dims already claimed by earlier mesh dims." + ) + if len(tgt_shard_dims) == 1: + tgt_shard_dim = tgt_shard_dims[0] else: - return Shard(shard_dim_map[p.dim]) - - output_placements = [ - _rewrite_shard_dim(p) if isinstance(p, Shard | _StridedShard) else p - for p in input_tgt_placements - ] - - return input_tgt_placements, output_placements + # Unflatten: one input dim maps to multiple output dims + # (e.g. (24,) → (2, 3, 4) gives 3 splits). Plain Shard + # always targets the split_id=0 output dim. + tgt_shard_dim = next( + ( + d + for d in tgt_shard_dims + if isinstance(self.rule[d], Split) + and cast(Split, self.rule[d]).split_id == 0 + ), + None, + ) + if tgt_shard_dim is None: + raise AssertionError( + f"No Split(split_id=0) found among unclaimed output dims " + f"{tgt_shard_dims} for Shard(dim={p.dim}) on mesh dim {mesh_dim}." + ) + cmd = self.rule[tgt_shard_dim] + if isinstance(cmd, Split) and isinstance(cmd.input_dim, Flatten): + first_dim = cmd.input_dim.input_dims[0] + if isinstance(first_dim, InputDim) and p.dim != first_dim.input_dim: + raise RuntimeError( + f"Shard(dim={p.dim}) through Split(Flatten(...), {cmd.group_shape}) " + f"is not supported yet for non-first flatten dims." + ) + if isinstance(cmd, (Split, InputDim)): + # Split/InputDim: 1:1 dim mapping, sharding transfers directly. + # Flatten needs stride computation below (multiple dims merge). + new_shapes = list(local_tensor_shapes) + new_shapes[p.dim] //= self.mesh_sizes[mesh_dim] + return Shard(tgt_shard_dim), new_shapes + if not isinstance(cmd, Flatten): + raise AssertionError(f"Expected Flatten, got {type(cmd)}") + first_dim = cmd.input_dims[0] + last_dim = cmd.input_dims[-1] + if not isinstance(first_dim, InputDim): + raise AssertionError(f"Expected InputDim, got {type(first_dim)}") + if not isinstance(last_dim, InputDim): + raise AssertionError(f"Expected InputDim, got {type(last_dim)}") + input_start_idx = first_dim.input_dim + if p.dim == input_start_idx: + output_placement: Placement = Shard(tgt_shard_dim) + else: + split_factor = math.prod(local_tensor_shapes[input_start_idx : p.dim]) + output_placement = _StridedShard(tgt_shard_dim, split_factor=split_factor) + # Uneven sharding on a non-last flatten dim breaks _StridedShard: + # split_factor (number of groups) must be the same on all devices, + # but uneven division of a non-last dim makes group count vary. + # E.g. [3,4]→[12] Shard(0) mesh=2: device 0 has 2 groups of 4, + # device 1 has 1 group of 4 — no consistent split_factor. + # The last dim is exempt: only group *size* varies, not count. + flatten_end = last_dim.input_dim + 1 + if local_tensor_shapes[p.dim] % self.mesh_sizes[ + mesh_dim + ] != 0 and not self._is_last_shard_in_flatten_range( + mesh_dim, placements, input_start_idx, flatten_end + ): + raise RuntimeError( + f"Cannot shard unevenly distributed tensor: " + f"dimension {p.dim} (size {local_tensor_shapes[p.dim]}) " + f"is not evenly divisible by mesh dimension " + f"{mesh_dim} (size {self.mesh_sizes[mesh_dim]}). " + f"Please redistribute the tensor before this operation." + ) + new_shapes = list(local_tensor_shapes) + new_shapes[p.dim] //= self.mesh_sizes[mesh_dim] + return output_placement, new_shapes + + def _rewrite_strided_shard( + self, + p: _StridedShard, + mesh_dim: int, + placements: Sequence[Placement], + strided_shard_claimed_dims: set[ClaimedDim], + local_tensor_shapes: list[int], + input_to_output_tensor_dims: dict[int, list[int]], + ) -> tuple[Placement, list[int]]: + """Rewrite _StridedShard placement to target the correct output dim. + + _StridedShard inputs arise from a prior flatten on a non-first dim + (produced by _rewrite_plain_shard above). The interesting case is + unflatten (Split rule): the split_factor may resolve to contiguous + sharding (producing Shard) or stay as _StridedShard. For + identity/flatten rules, falls through to the fallback and keeps the + placement as-is. + + Returns the output placement and a new local_tensor_shapes with this + mesh dim's division applied. + """ + tgt_shard_dims = [ + d + for d in input_to_output_tensor_dims[p.dim] + if ClaimedDim(p.dim, d) not in strided_shard_claimed_dims + ] + # Phase 1: resolve SS → Shard. If an output dim's Split has a + # group_shape prefix matching the split_factor, the strided pattern + # is fully captured by the Split, so SS simplifies to Shard. + # E.g. unflatten (6, 4) → (2, 3, 4) with SS(0, sf=2) on mesh (3): + # sf=2 means 2 groups of contiguous data in dim 0. Split into + # (2, 3, 4) gives group_shape=(2, 3); prod(group_shape[:1])=2==sf, + # so the strided pattern lands exactly on output dim 1 → Shard(1). + for candidate_dim in tgt_shard_dims: + cmd = self.rule[candidate_dim] + if isinstance(cmd, Split): + expected_sf = self._expected_split_factor( + cmd, p.dim, mesh_dim, placements + ) + if expected_sf != p.split_factor: + continue + strided_shard_claimed_dims.add(ClaimedDim(p.dim, candidate_dim)) + new_shapes = list(local_tensor_shapes) + new_shapes[p.dim] //= self.mesh_sizes[mesh_dim] + return Shard(candidate_dim), new_shapes + + # Phase 2: keep SS as SS. Phase 1 is tried first because we prefer + # resolving to the simpler Shard when possible. + tgt_shard_dim = self._find_keep_ss_dim(tgt_shard_dims, p, mesh_dim) + + if tgt_shard_dim is None: + if self.strict_view and any( + isinstance(self.rule[d], Split) for d in tgt_shard_dims + ): + raise RuntimeError( + f"Cannot unflatten tensor with _StridedShard placement: " + f"split_factor={p.split_factor} does not match any output " + f"dimension. This typically means the _StridedShard placement " + f"was constructed with a split_factor that is incompatible " + f"with the unflatten shape. Please redistribute the tensor " + f"before this operation." + ) + if len(tgt_shard_dims) == 0: + raise AssertionError( + f"No unclaimed output dims for _StridedShard(dim={p.dim}) " + f"on mesh dim {mesh_dim}." + ) + # Fallback for identity/flatten: tgt_shard_dims has exactly one + # element, so [0] is correct. For Split rules this is unreachable + # in practice — the analysis phase rejects mismatched split_factors + # via shard_allowed, forcing redistribution before we get here. + tgt_shard_dim = tgt_shard_dims[0] + new_shapes = list(local_tensor_shapes) + new_shapes[p.dim] //= self.mesh_sizes[mesh_dim] + return _StridedShard(tgt_shard_dim, split_factor=p.split_factor), new_shapes def register_op_strategy_map( @@ -783,12 +1291,17 @@ def reshape_strategy(op_schema: OpSchema) -> StrategyType: placements=tuple(input_tgt_placements), mesh=mesh, tensor_meta=input_src_spec.tensor_meta, + use_strided_shard_as_shard_order=False, ) redistribute_costs: list[list[float]] = [ generate_redistribute_costs(input_strategy, input_tgt_spec) ] - output_spec = DTensorSpec(mesh=mesh, placements=tuple(output_placements)) + output_spec = DTensorSpec( + mesh=mesh, + placements=tuple(output_placements), + use_strided_shard_as_shard_order=False, + ) output_strategy.strategies.append( OpSpec( output_specs=output_spec, @@ -801,6 +1314,7 @@ def reshape_strategy(op_schema: OpSchema) -> StrategyType: register_op_strategy_map(aten.squeeze.default, torch.squeeze) +register_op_strategy_map(aten.squeeze_.default, torch.squeeze) register_op_strategy_map( aten.squeeze_.dim, torch.squeeze, schema_info=RuntimeSchemaInfo(1) ) diff --git a/torch/distributed/tensor/_ops/single_dim_strategy.py b/torch/distributed/tensor/_ops/single_dim_strategy.py index c64b89a7c156f..722cb8c24a8a2 100644 --- a/torch/distributed/tensor/_ops/single_dim_strategy.py +++ b/torch/distributed/tensor/_ops/single_dim_strategy.py @@ -23,7 +23,6 @@ OpSchema, OpSpec, OpStrategy, - PlacementList, RuntimeSchemaInfo, StrategyType, TupleStrategy, @@ -76,6 +75,10 @@ class _SingleDimStrategyInfo: func: _SingleDimStrategyFunc allow_unbacked_sharding: bool | None = field(default=None) allow_uneven_sharding: bool = field(default=False) + # Positions (in args_schema) of args that may live on a different mesh + # than the op's compute mesh. These args must be Replicate. + # See Note [Multi-mesh args] in expand_to_full_mesh_op_strategy. + different_mesh_args: list[int] | None = field(default=None) # Delegate to func so this can be used interchangeably with a raw # _SingleDimStrategyFunc (e.g. in tests that call strategy functions directly). @@ -85,29 +88,39 @@ def __call__(self, *args, **kwargs): def _insert_single_dim_replication_strategy( single_dim_strategies_with_placeholders: list[ - list[Placement | _ShardingPlaceholder] + list[Placement | _ShardingPlaceholder | None] ], num_outputs: int, num_input_tensors: int, -) -> list[list[Placement | _ShardingPlaceholder]]: + output_tensor_meta: TensorMeta | Sequence[TensorMeta | None] | None = None, +) -> list[list[Placement | _ShardingPlaceholder | None]]: """ Inserts the [Replicate(), Replicate(), ...] strategy after asserting that such strategy does not yet exist. + For ops with masked-off outputs (e.g. backward ops with output_mask), output positions + where output_tensor_meta is None are set to None in the all-Replicate rule. """ for strategy in single_dim_strategies_with_placeholders: - if all(isinstance(p, Replicate) for p in strategy): - raise AssertionError - single_dim_strategies_with_placeholders.insert( - 0, [Replicate()] * (num_outputs + num_input_tensors) - ) + if all(isinstance(p, Replicate) or p is None for p in strategy): + return single_dim_strategies_with_placeholders + total_len = num_outputs + num_input_tensors + replicate_rule: list[Placement | _ShardingPlaceholder | None] = [ + Replicate() + ] * total_len + # Set None for masked-off output positions based on output_tensor_meta + if isinstance(output_tensor_meta, Sequence): + for i, meta in enumerate(output_tensor_meta): + if meta is None and i < num_outputs: + replicate_rule[i] = None + single_dim_strategies_with_placeholders.insert(0, replicate_rule) return single_dim_strategies_with_placeholders def _fill_single_dim_strategy_placeholders( unique_input_placements: set[Placement], single_dim_strategies_with_placeholders: list[ - list[Placement | _ShardingPlaceholder] + list[Placement | _ShardingPlaceholder | None] ], -) -> list[list[Placement]]: +) -> list[list[Placement | None]]: """ Replace any _ShardingPlaceholder with the specific Sharding types used by the inputs in op_schema. Supports implicit replication. @@ -137,25 +150,29 @@ def _fill_single_dim_strategy_placeholders( # if any of the placements is a placeholder, we need to expand the strategy # to all possible combinations of placements - expanded_strategies_over_one_mesh_dim: list[list[Placement]] = [] + expanded_strategies_over_one_mesh_dim: list[list[Placement | None]] = [] for s in single_dim_strategies_with_placeholders: if any(isinstance(p, _ShardingPlaceholder) for p in s): for shard_builder in shard_builders.values(): - expanded_strategy: list[Placement] = [] + expanded_strategy: list[Placement | None] = [] for maybe_placeholder in s: if isinstance(maybe_placeholder, _ShardingPlaceholder): # we combine the tensor dim to shard from the placeholder # with other metadata (e.g. split_factor) from the sharding class expanded_strategy.append(shard_builder(maybe_placeholder.dim)) + elif maybe_placeholder is None: + expanded_strategy.append(None) else: if not isinstance(maybe_placeholder, Placement): raise AssertionError expanded_strategy.append(maybe_placeholder) expanded_strategies_over_one_mesh_dim.append(expanded_strategy) else: - if not all(isinstance(p, Placement) for p in s): + if not all(isinstance(p, Placement) or p is None for p in s): raise AssertionError - expanded_strategies_over_one_mesh_dim.append(cast(list[Placement], (s))) + expanded_strategies_over_one_mesh_dim.append( + cast(list[Placement | None], (s)) + ) return expanded_strategies_over_one_mesh_dim @@ -193,7 +210,7 @@ def _count(obj: Any) -> int: if isinstance(obj, OpStrategy): return 1 elif isinstance(obj, TupleStrategy): - return len(obj.children) + return sum(1 for child in obj.children if child is not None) elif isinstance(obj, (list, tuple)): return sum(_count(child) for child in obj) return 0 @@ -208,14 +225,15 @@ def _count(obj: Any) -> int: def _build_output_specs( mesh: DeviceMesh, - per_mesh_dim_placements: list[tuple[Placement, ...]], + per_mesh_dim_placements: list[tuple[Placement | None, ...]], num_outputs: int, output_metas: tuple[TensorMeta | None, ...], ) -> DTensorSpec | tuple[DTensorSpec | None, ...]: """Build output spec(s) by transposing per-mesh-dim placements to per-output. per_mesh_dim_placements is indexed [mesh_dim][output_idx]. output_metas must - have exactly num_outputs elements. + have exactly num_outputs elements. Outputs where output_metas[i] is None + (masked-off outputs) produce None specs. """ if num_outputs <= 0: raise AssertionError(f"Expected num_outputs > 0, got {num_outputs}") @@ -224,16 +242,21 @@ def _build_output_specs( f"Expected {num_outputs} output_metas, got {len(output_metas)}" ) - def _placements_for_output(out_idx: int) -> tuple[Placement, ...]: - return tuple(out[out_idx] for out in per_mesh_dim_placements) + def _spec_for_output(out_idx: int) -> DTensorSpec | None: + if output_metas[out_idx] is None: + return None + placements = tuple( + cast(Placement, out[out_idx]) for out in per_mesh_dim_placements + ) + return DTensorSpec(mesh, placements, tensor_meta=output_metas[out_idx]) if num_outputs > 1: - return tuple( - DTensorSpec(mesh, _placements_for_output(i), tensor_meta=output_metas[i]) - for i in range(num_outputs) - ) + return tuple(_spec_for_output(i) for i in range(num_outputs)) else: - return DTensorSpec(mesh, _placements_for_output(0), tensor_meta=output_metas[0]) + spec = _spec_for_output(0) + if spec is None: + raise AssertionError("Single-output op cannot have None output meta") + return spec class _PreparedSingleDimStrategy: @@ -249,8 +272,8 @@ class _PreparedSingleDimStrategy: allowed_partial_per_input for graph search neighbor generation. """ - strategy_lookup: dict[tuple[Placement, ...], tuple[Placement, ...]] - expanded_strategies: list[list[Placement]] + strategy_lookup: dict[tuple[Placement | None, ...], tuple[Placement | None, ...]] + expanded_strategies: list[list[Placement | None]] num_outputs: int num_inputs: int output_metas: tuple[TensorMeta | None, ...] @@ -282,18 +305,63 @@ def __init__( if isinstance(strategy_fn, _SingleDimStrategyInfo): self.allow_unbacked_sharding = strategy_fn.allow_unbacked_sharding self.allow_uneven_sharding = strategy_fn.allow_uneven_sharding + different_mesh_args = strategy_fn.different_mesh_args func = strategy_fn.func else: self.allow_unbacked_sharding = None self.allow_uneven_sharding = False + different_mesh_args = None func = strategy_fn + # Determine element_mesh from the first OpStrategy arg. For foreach + # per-element schemas the element's inputs may live on a smaller + # sub-mesh than the global compute_mesh. + self.element_mesh: DeviceMesh | None = None + for arg in op_schema.args_schema: + if isinstance(arg, OpStrategy): + self.element_mesh = arg.strategies[0].output_spec.mesh + break + + # Validate that all inputs are on the same mesh (except + # different_mesh_args which are explicitly allowed to differ). + if self.element_mesh is not None: + allowed = set(different_mesh_args or []) + for i, arg in enumerate(op_schema.args_schema): + if isinstance(arg, OpStrategy) and i not in allowed: + arg_mesh = arg.strategies[0].output_spec.mesh + if arg_mesh != self.element_mesh: + raise ValueError( + f"Cannot run {op_schema.op} on inputs with different " + f"meshes: got {self.element_mesh} and {arg_mesh}" + ) + + # Remap different_mesh_args from args_schema positions to + # OpStrategy-only positions. Non-OpStrategy args (e.g. empty lists) + # are filtered out by expand_to_full_mesh_op_strategy, shifting later + # indices. + self.remapped_different_mesh_args: list[int] | None = None + if different_mesh_args is not None: + schema_to_strategy: dict[int, int] = {} + strategy_pos = 0 + for schema_pos, arg in enumerate(op_schema.args_schema): + if isinstance(arg, OpStrategy): + schema_to_strategy[schema_pos] = strategy_pos + strategy_pos += 1 + self.remapped_different_mesh_args = [ + schema_to_strategy[i] + for i in different_mesh_args + if i in schema_to_strategy + ] + if num_inputs is None: num_inputs = _get_num_tensor_inputs(op_schema) self.num_inputs = num_inputs - strategies_with_placeholders = func( - op_schema.op, op_schema.args_meta, op_schema.kwargs_meta + # Strategy functions may return None in output positions for masked-off + # outputs (e.g. backward ops with output_mask). Widen the type here. + strategies_with_placeholders = cast( + list[list[Placement | _ShardingPlaceholder | None]], + func(op_schema.op, op_schema.args_meta, op_schema.kwargs_meta), ) # Validate strategy length against the op schema. The schema is the @@ -327,7 +395,10 @@ def __init__( self.num_outputs = num_outputs strategies_with_placeholders = _insert_single_dim_replication_strategy( - strategies_with_placeholders, num_outputs, num_inputs + strategies_with_placeholders, + num_outputs, + num_inputs, + output_tensor_meta, ) unique_input_placements = _get_unique_placements(op_schema) @@ -350,6 +421,8 @@ def __init__( for strategy in self.expanded_strategies: for input_idx in range(num_inputs): p = strategy[num_outputs + input_idx] + if p is None: + continue if _is_sharding(p): self.allowed_sharding_per_input[input_idx].add(p) elif isinstance(p, Partial): @@ -378,7 +451,7 @@ def try_propagate( """ from torch.distributed.tensor._ops.utils import is_tensor_shardable - selected_output_placements: list[tuple[Placement, ...]] = [] + selected_output_placements: list[tuple[Placement | None, ...]] = [] for mesh_dim in range(mesh.ndim): input_placements_for_dim = tuple( placements[mesh_dim] for placements in input_placements @@ -464,15 +537,18 @@ def expanded_strategy( base_name = op_name.split("::")[1].split(".")[0] is_inplace = base_name.endswith("_") + element_mesh = prepared_strategy.element_mesh or mesh + return expand_to_full_mesh_op_strategy( - mesh, + element_mesh, op_schema, - cast(list[PlacementList], prepared_strategy.expanded_strategies), + prepared_strategy.expanded_strategies, output_tensor_meta=output_tensor_meta, inplace_op=is_inplace, input_index=prepared_strategy.num_outputs, allow_unbacked_sharding=prepared_strategy.allow_unbacked_sharding, allow_uneven_sharding=prepared_strategy.allow_uneven_sharding, + different_mesh_args=prepared_strategy.remapped_different_mesh_args, ) return expanded_strategy @@ -494,12 +570,14 @@ def _create_expanded_strategy( # Unhashable types (SymInts), skip caching return _create_expanded_strategy_impl(op_schema, output_tensor_meta) - def _translate_foreach_op_schema( - op_schema: OpSchema, output_tensor_meta: Sequence[TensorMeta], index: int - ) -> tuple[OpSchema, TensorMeta]: - """Translate foreach op to per-element version of schema.""" + def _translate_list_op_schema( + op_schema: OpSchema, + output_tensor_meta: Sequence[TensorMeta] | None, + index: int, + ) -> tuple[OpSchema, TensorMeta | None]: + """Translate foreach/fused op to per-element version of schema.""" op_parts = str(op_schema.op).split(".") - base_op_name = op_parts[-2].replace("_foreach_", "") + op_name = op_parts[-2] foreach_variant = op_parts[-1] # select per-element inputs, outputs @@ -509,7 +587,30 @@ def _translate_foreach_op_schema( (op_schema.args_schema, op_schema.kwargs_schema), is_leaf=lambda x: isinstance(x, TupleStrategy), ) - target_output_meta = output_tensor_meta[index] + # For inplace ops, output_tensor_meta is None + target_output_meta = ( + output_tensor_meta[index] if output_tensor_meta is not None else None + ) + + # Strip the prefix to get the base op name and find the per-element op. + # Fused ops (e.g. _fused_adam) have no per-element ATen equivalent, + # so we keep the original op unchanged. + if op_name.startswith("_foreach_"): + base_op_name = op_name.replace("_foreach_", "", 1) + elif op_name.startswith("_amp_foreach_"): + base_op_name = op_name.replace("_amp_foreach_", "", 1) + else: + # Fused ops or unknown: keep original op, no translation + target_op = op_schema.op + op_schema = OpSchema( + target_op, # type: ignore[arg-type] + args_schema=tuple(target_args), + kwargs_schema=op_schema.kwargs_schema, + ) + return op_schema, target_output_meta + + # Strip trailing underscore for inplace ops + base_op_name = base_op_name.removesuffix("_") # figure out target op variant variant_map = { @@ -558,7 +659,7 @@ def expanded_foreach_strategy( child_strategies: list[StrategyType] = [] for tensorlist_i in range(tensorlist_len): - per_index_schema, per_index_output_meta = _translate_foreach_op_schema( + per_index_schema, per_index_output_meta = _translate_list_op_schema( op_schema, output_tensor_meta, # type: ignore[arg-type] tensorlist_i, @@ -575,7 +676,16 @@ def expanded_foreach_strategy( return TupleStrategy(children=child_strategies) # TODO maybe this could be helped by adding a new 'tag' to the OpOverload? - if op_schema.op.name().startswith("aten::_foreach_"): + # Only use the foreach path if the op has TupleStrategy inputs (i.e., actual + # list-of-tensor args). The name prefix alone is insufficient because ops like + # _fused_rms_norm share the "_fused_" prefix but are not foreach/fused-optimizer ops. + op_name = op_schema.op.name() + has_tuple_strategy = any( + isinstance(arg, TupleStrategy) for arg in op_schema.args_schema + ) + if has_tuple_strategy and op_name.startswith( + ("aten::_foreach_", "aten::_amp_foreach_", "aten::_fused_") + ): return expanded_foreach_strategy return _create_expanded_strategy(op_schema, output_tensor_meta) @@ -586,6 +696,7 @@ def register_single_dim_strategy( schema_info: RuntimeSchemaInfo | None = None, allow_unbacked_sharding: bool | None = None, allow_uneven_sharding: bool = False, + different_mesh_args: list[int] | None = None, ) -> Callable[[_SingleDimStrategyFunc], _SingleDimStrategyFunc]: """ Registers a single_dim_strategy function for the given op. @@ -634,6 +745,7 @@ def wrapper(impl): func=impl, allow_unbacked_sharding=allow_unbacked_sharding, allow_uneven_sharding=allow_uneven_sharding, + different_mesh_args=different_mesh_args, ) registration_wrapper(info) return impl @@ -817,6 +929,10 @@ def _dijkstra_expand_single_dim_strategy_to_mesh( total_bytes = spec.tensor_meta.dtype.itemsize * math.prod( spec.tensor_meta.shape ) + # TODO: is_shard() misses _StridedShard, use spec.num_shards instead. + # Not fixing yet: the overestimate biases Dijkstra toward redistributing + # away from _StridedShard, which is the safer default until _StridedShard + # is fully validated. num_shards = 1 for i, p in enumerate(spec.placements): if p.is_shard(): diff --git a/torch/distributed/tensor/_ops/strategy_validation.py b/torch/distributed/tensor/_ops/strategy_validation.py index ad2bf360994db..3d94de2d1455d 100644 --- a/torch/distributed/tensor/_ops/strategy_validation.py +++ b/torch/distributed/tensor/_ops/strategy_validation.py @@ -62,6 +62,7 @@ SKIP_OPS: dict[str, str] = { "bernoulli": "non-deterministic (random sampling)", "empty_like": "uninitialized memory", + "exponential": "non-deterministic (random sampling)", "new_empty": "uninitialized memory", "new_empty_strided": "uninitialized memory", "nn.functional.dropout": "non-deterministic (random masking)", @@ -158,6 +159,7 @@ def is_fully_replicated(placements: tuple[Placement, ...]) -> bool: def is_trivial_shard(p: Placement, tensor_shape: tuple[int, ...]) -> bool: """Check if placement is a Shard on a size-1 dimension.""" + # NOTE: isinstance(_, Shard) does not match _StridedShard; see _is_shard_like(). return ( isinstance(p, Shard) and p.dim < len(tensor_shape) and tensor_shape[p.dim] == 1 ) @@ -429,6 +431,121 @@ def _create_partial_input( return LocalTensor(local_tensors) +def _shard_tensors( + tensors: list[tuple[str, torch.Tensor]], + input_placements: tuple[Placement, ...], + world_size: int, + mesh: DeviceMesh, + mask_shift: int = 0, +) -> list[LocalTensor | torch.Tensor]: + """Create sharded LocalTensors from tensors according to placements.""" + local_tensors: list[LocalTensor | torch.Tensor] = [] + for tensor_idx, ((name, tensor), placement) in enumerate( + zip(tensors, input_placements) + ): + if isinstance(placement, Partial): + local_tensor = _create_partial_input( + tensor, placement, world_size, tensor_idx, mask_shift + ) + elif isinstance(placement, Replicate): + _tmp = {r: tensor.clone() for r in range(world_size)} + # pyrefly: ignore [bad-argument-type, bad-argument-count] + local_tensor = LocalTensor(_tmp) + elif isinstance(placement, Shard): + shard_dim = placement.dim + chunks = tensor.tensor_split(world_size, dim=shard_dim) + _tmp = { + r: chunks[r].clone(memory_format=torch.contiguous_format) + for r in range(world_size) + } + # pyrefly: ignore [bad-argument-type, bad-argument-count] + local_tensor = LocalTensor(_tmp) + else: + dt = distribute_tensor(tensor.clone(), mesh, (placement,)) + local_tensor = dt.to_local() + local_tensors.append(local_tensor) + return local_tensors + + +def _compare_outputs( + local_output: Any, + ground_truth: torch.Tensor | list[torch.Tensor], + output_placements: tuple[Placement, ...], + mesh: DeviceMesh, + world_size: int, +) -> tuple[bool, str]: + """Compare op output (wrapped as DTensor) against ground truth.""" + if isinstance(local_output, (list, tuple)): + local_outputs = list(local_output) + else: + local_outputs = [local_output] + + if isinstance(ground_truth, list): + ground_truths = ground_truth + else: + ground_truths = [ground_truth] + + if len(local_outputs) != len(ground_truths): + return ( + False, + f"Output count mismatch: got {len(local_outputs)}, " + f"expected {len(ground_truths)}", + ) + + if len(local_outputs) != len(output_placements): + return ( + False, + f"Output count mismatch with placements: " + f"got {len(local_outputs)}, expected {len(output_placements)}", + ) + + for i, (local_out, gt, out_plc) in enumerate( + zip(local_outputs, ground_truths, output_placements) + ): + if not isinstance(local_out, torch.Tensor): + return False, f"Local output[{i}] is not a tensor: {type(local_out)}" + + if not isinstance(local_out, LocalTensor): + return False, f"LocalTensor inputs produced non-LocalTensor output[{i}]" + + output_dt = DTensor.from_local( + local_out, + mesh, + (out_plc,), + shape=gt.shape, + stride=gt.stride(), + ) + + if isinstance(out_plc, Replicate): + local_values = [local_out._local_tensors[r] for r in range(world_size)] + all_same = all( + torch.allclose(local_values[0], lv, atol=1e-5, rtol=1e-5) + for lv in local_values[1:] + ) + if not all_same: + return ( + False, + f"Replicate output[{i}] but local values differ across ranks", + ) + + full_output = output_dt.redistribute(mesh, (Replicate(),)).to_local() + + if isinstance(full_output, LocalTensor): + full_output = full_output._local_tensors[0] + + if gt.shape != full_output.shape: + return ( + False, + f"Shape mismatch[{i}]: expected {gt.shape}, got {full_output.shape}", + ) + + if not torch.allclose(gt, full_output, atol=1e-5, rtol=1e-5, equal_nan=True): + max_diff = (gt - full_output).abs().max().item() + return False, f"Value mismatch[{i}]: max_diff={max_diff:.6f}" + + return True, "" + + def validate_combination( op: Callable[..., Any], sample_input: SampleInput, @@ -438,10 +555,13 @@ def validate_combination( world_size: int = 2, mesh: DeviceMesh | None = None, mask_shift: int = 0, -) -> tuple[bool, str]: +) -> tuple[bool | None, str]: """ Validate a single placement combination against ground truth. + Returns (True, "") if valid, (False, error_msg) if invalid, or + (None, reason) if the combination cannot be tested (e.g. uneven shards). + The validation logic: 1. Shard inputs according to input placements to get local tensors 2. Run the raw op on local tensors (bypassing DTensor dispatch) @@ -467,33 +587,17 @@ def validate_combination( device = tensors[0][1].device.type if tensors else "cpu" mesh = init_device_mesh(device, (world_size,)) - local_tensors = [] - for tensor_idx, ((name, tensor), placement) in enumerate( - zip(tensors, combination[0]) - ): - if isinstance(placement, Partial): - local_tensor = _create_partial_input( - tensor, placement, world_size, tensor_idx, mask_shift - ) - elif isinstance(placement, Replicate): - _tmp = {r: tensor.clone() for r in range(world_size)} - # pyrefly: ignore [bad-argument-type, bad-argument-count] - local_tensor = LocalTensor(_tmp) - elif isinstance(placement, Shard): - # Create sharded LocalTensor directly to work in LocalTensorMode - shard_dim = placement.dim - chunks = tensor.tensor_split(world_size, dim=shard_dim) - _tmp = { - r: chunks[r].clone(memory_format=torch.contiguous_format) - for r in range(world_size) - } - # pyrefly: ignore [bad-argument-type, bad-argument-count] - local_tensor = LocalTensor(_tmp) - else: - # Fallback for other placement types - dt = distribute_tensor(tensor.clone(), mesh, (placement,)) - local_tensor = dt.to_local() - local_tensors.append(local_tensor) + # Uneven shards produce SymInt in LocalTensor's wrapper shape, + # which breaks C++ overload resolution before __torch_dispatch__ + # can intercept. Return None to signal "untestable". + for (name, tensor), placement in zip(tensors, combination[0]): + if isinstance(placement, Shard): + if tensor.size(placement.dim) % world_size != 0: + return None, "uneven shard" + + local_tensors = _shard_tensors( + tensors, combination[0], world_size, mesh, mask_shift + ) local_idx = 0 @@ -515,84 +619,95 @@ def _replace_with_local(a): local_output = op(local_input, *local_args, **local_kwargs) - # Normalize to list for uniform handling of single/multi-output ops - if isinstance(local_output, (list, tuple)): - local_outputs = list(local_output) - else: - local_outputs = [local_output] + return _compare_outputs( + local_output, ground_truth, combination[1], mesh, world_size + ) - if isinstance(ground_truth, list): - ground_truths = ground_truth - else: - ground_truths = [ground_truth] + except Exception as e: + # TODO: This is too broad. Consider: (1) explicit checks for shard dim + # validity and shape compatibility before calling tensor_split/from_local, + # (2) scoped try/except around op() and redistribute() that raise specific + # exceptions (e.g., UnsupportedRedistribute, OpError), and (3) only + # catching those here, letting real bugs propagate. + return False, f"Exception: {type(e).__name__}: {e}" - if len(local_outputs) != len(ground_truths): - return ( - False, - f"Output count mismatch: got {len(local_outputs)}, expected {len(ground_truths)}", - ) - if len(local_outputs) != len(combination[1]): - return ( - False, - f"Output count mismatch with placements: " - f"got {len(local_outputs)}, expected {len(combination[1])}", - ) +def extract_tensors_from_args( + args: tuple[Any, ...], + kwargs: dict[str, Any], +) -> list[tuple[str, torch.Tensor]]: + """Extract tensor arguments from captured aten args/kwargs. - for i, (local_out, gt, out_plc) in enumerate( - zip(local_outputs, ground_truths, combination[1]) - ): - if not isinstance(local_out, torch.Tensor): - return False, f"Local output[{i}] is not a tensor: {type(local_out)}" + Unlike extract_tensors_from_sample which walks SampleInput pytrees, + this walks the flat aten-level args and kwargs directly. + """ + tensors: list[tuple[str, torch.Tensor]] = [] + idx = 0 - if not isinstance(local_out, LocalTensor): - return False, f"LocalTensor inputs produced non-LocalTensor output[{i}]" + def _collect(x: Any) -> Any: + nonlocal idx + if isinstance(x, torch.Tensor): + tensors.append((f"tensor_{idx}", x)) + idx += 1 + return x - output_dt = DTensor.from_local( - local_out, - mesh, - (out_plc,), - shape=gt.shape, - stride=gt.stride(), - ) + pytree.tree_map(_collect, args) + pytree.tree_map(_collect, kwargs) + return tensors - if isinstance(out_plc, Replicate): - local_values = [local_out._local_tensors[r] for r in range(world_size)] - all_same = all( - torch.allclose(local_values[0], lv, atol=1e-3, rtol=1e-3) - for lv in local_values[1:] - ) - if not all_same: - return ( - False, - f"Replicate output[{i}] but local values differ across ranks", - ) - full_output = output_dt.redistribute(mesh, (Replicate(),)).to_local() +def validate_aten_combination( + aten_op: OpOverload, + captured_args: tuple[Any, ...], + captured_kwargs: dict[str, Any], + ground_truth: torch.Tensor | list[torch.Tensor], + combination: PlacementCombination, + world_size: int, + mesh: DeviceMesh, + mask_shift: int = 0, +) -> tuple[bool | None, str]: + """Validate a placement combination using aten-level captured args. + + Works directly with aten op args/kwargs instead of SampleInput pytrees. + Replaces tensors in the flat args/kwargs with sharded LocalTensors, + calls the aten op, and compares output. - if isinstance(full_output, LocalTensor): - full_output = full_output._local_tensors[0] + Returns (True, ""), (False, error_msg), or (None, reason) if untestable. + """ + try: + tensors = extract_tensors_from_args(captured_args, captured_kwargs) + if not tensors: + return False, "No tensor args in captured aten call" - if gt.shape != full_output.shape: - return ( - False, - f"Shape mismatch[{i}]: expected {gt.shape}, got {full_output.shape}", - ) + for (name, tensor), placement in zip(tensors, combination[0]): + if isinstance(placement, Shard): + if tensor.size(placement.dim) % world_size != 0: + return None, "uneven shard" + + local_tensors = _shard_tensors( + tensors, combination[0], world_size, mesh, mask_shift + ) + + local_idx = 0 + + def _replace_with_local(a: Any) -> Any: + nonlocal local_idx + if isinstance(a, torch.Tensor): + local = local_tensors[local_idx] + local_idx += 1 + return local + return a - if not torch.allclose( - gt, full_output, atol=1e-3, rtol=1e-3, equal_nan=True - ): - max_diff = (gt - full_output).abs().max().item() - return False, f"Value mismatch[{i}]: max_diff={max_diff:.6f}" + local_args = pytree.tree_map(_replace_with_local, captured_args) + local_kwargs = pytree.tree_map(_replace_with_local, captured_kwargs) - return True, "" + local_output = aten_op(*local_args, **local_kwargs) + + return _compare_outputs( + local_output, ground_truth, combination[1], mesh, world_size + ) except Exception as e: - # TODO: This is too broad. Consider: (1) explicit checks for shard dim - # validity and shape compatibility before calling tensor_split/from_local, - # (2) scoped try/except around op() and redistribute() that raise specific - # exceptions (e.g., UnsupportedRedistribute, OpError), and (3) only - # catching those here, letting real bugs propagate. return False, f"Exception: {type(e).__name__}: {e}" @@ -695,34 +810,41 @@ def _extract_rules_from_op_strategy( class _CaptureAtenOp(torch.utils._python_dispatch.TorchDispatchMode): - """Dispatch mode that captures aten ops called and their args.""" + """Dispatch mode that captures aten ops called, their args, and return values.""" def __init__(self, target_op_name: str = ""): self.target_op_name = target_op_name.lower() - self.all_ops: list[tuple[OpOverload, tuple[Any, ...], dict[str, Any]]] = [] + self.all_ops: list[tuple[OpOverload, tuple[Any, ...], dict[str, Any], Any]] = [] self.best_match: OpOverload | None = None self.best_match_args: tuple[Any, ...] | None = None self.best_match_kwargs: dict[str, Any] | None = None + self.best_match_result: Any = None def __torch_dispatch__(self, func, types, args=(), kwargs=None): if kwargs is None: kwargs = {} + result = func(*args, **kwargs) if func.namespace == "aten": - self.all_ops.append((func, args, kwargs)) + self.all_ops.append((func, args, kwargs, result)) op_name = func.name().split("::")[1].split(".")[0].lower() if self.target_op_name and self.target_op_name in op_name: if self.best_match is None: self.best_match = func self.best_match_args = args self.best_match_kwargs = kwargs - return func(*args, **kwargs) + self.best_match_result = result + return result def get_aten_op_for_sample( op: Callable[..., Any], sample: SampleInput, op_name: str = "" -) -> tuple[OpOverload | None, tuple[Any, ...], dict[str, Any]]: +) -> _CaptureAtenOp: """ - Determine the actual aten op that will be dispatched for a given sample. + Capture aten ops dispatched for a given sample. + + Returns the _CaptureAtenOp object containing all captured ops with their + args, kwargs, and return values. Use best_match for the primary op or + all_ops for exhaustive iteration. """ with _CaptureAtenOp(op_name) as capture: try: @@ -733,16 +855,15 @@ def get_aten_op_for_sample( except Exception: pass - if capture.best_match is not None: - captured_op = capture.best_match - captured_args = capture.best_match_args - captured_kwargs = capture.best_match_kwargs - elif capture.all_ops: - captured_op, captured_args, captured_kwargs = capture.all_ops[0] - else: - return None, (), {} + # Populate best_match from first op if target match wasn't found + if capture.best_match is None and capture.all_ops: + first_op, first_args, first_kwargs, first_result = capture.all_ops[0] + capture.best_match = first_op + capture.best_match_args = first_args + capture.best_match_kwargs = first_kwargs + capture.best_match_result = first_result - return captured_op, captured_args, captured_kwargs + return capture def query_single_dim_strategy( @@ -1091,12 +1212,17 @@ def _validate_with_mitigations( world_size: int, mesh: DeviceMesh, mitigations: _FalsePositiveMitigations, -) -> bool: - """Validate a combination, including false positive mitigation re-checks.""" +) -> bool | None: + """Validate a combination, including false positive mitigation re-checks. + + Returns True (valid), False (invalid), or None (untestable). + """ combo: PlacementCombination = (input_placements, output_placements) is_valid, _ = validate_combination( op, sample, tensors, combo, ground_truth, world_size, mesh ) + if is_valid is None: + return None # Flipped-mask mitigation: the checkerboard mask that controls offset # signs (for P(sum)/P(avg)) or rank ownership (for P(min)/P(max)) is @@ -1174,6 +1300,115 @@ def _validate_with_mitigations( return is_valid +@dataclass +class _AtenFalsePositiveMitigations: + """Bundle of negated variants for aten-level false positive detection. + + Unlike _FalsePositiveMitigations, this works with captured aten args/kwargs + directly, without SampleInput. The rounding_mode mitigation is skipped + since it's an OpInfo-level concept that doesn't appear in captured aten kwargs. + """ + + negated_args: tuple[Any, ...] | None = None + negated_kwargs: dict[str, Any] | None = None + negated_ground_truth: torch.Tensor | list[torch.Tensor] | None = None + + +def _negate_tensors_in_tree(tree: Any) -> Any: + """Negate all tensors in a pytree structure.""" + + def _negate(x: Any) -> Any: + if isinstance(x, torch.Tensor): + return -x + return x + + return pytree.tree_map(_negate, tree) + + +def _prepare_aten_mitigations( + aten_op: OpOverload, + captured_args: tuple[Any, ...], + captured_kwargs: dict[str, Any], +) -> _AtenFalsePositiveMitigations: + """Create negated variants for aten-level false positive detection.""" + m = _AtenFalsePositiveMitigations() + try: + m.negated_args = _negate_tensors_in_tree(captured_args) + m.negated_kwargs = _negate_tensors_in_tree(captured_kwargs) + result = aten_op(*m.negated_args, **m.negated_kwargs) + if _is_tensor_output(result): + m.negated_ground_truth = _to_ground_truth(result) + else: + m.negated_args = None + m.negated_kwargs = None + except Exception: + m.negated_args = None + m.negated_kwargs = None + return m + + +def _validate_aten_with_mitigations( + aten_op: OpOverload, + captured_args: tuple[Any, ...], + captured_kwargs: dict[str, Any], + input_placements: tuple[Placement, ...], + output_placements: tuple[Placement, ...], + ground_truth: torch.Tensor | list[torch.Tensor], + world_size: int, + mesh: DeviceMesh, + mitigations: _AtenFalsePositiveMitigations, +) -> bool | None: + """Validate an aten-level combination with false positive mitigations. + + Returns True (valid), False (invalid), or None (untestable). + """ + combo: PlacementCombination = (input_placements, output_placements) + is_valid, _ = validate_aten_combination( + aten_op, + captured_args, + captured_kwargs, + ground_truth, + combo, + world_size, + mesh, + ) + if is_valid is None: + return None + + if is_valid and has_any_partial(input_placements, output_placements): + is_valid, _ = validate_aten_combination( + aten_op, + captured_args, + captured_kwargs, + ground_truth, + combo, + world_size, + mesh, + mask_shift=1, + ) + + if ( + is_valid + and mitigations.negated_args is not None + and has_pmin_pmax(input_placements, output_placements) + ): + if mitigations.negated_kwargs is None: + raise AssertionError("negated_kwargs must not be None") + if mitigations.negated_ground_truth is None: + raise AssertionError("negated_ground_truth must not be None") + is_valid, _ = validate_aten_combination( + aten_op, + mitigations.negated_args, + mitigations.negated_kwargs, + mitigations.negated_ground_truth, + combo, + world_size, + mesh, + ) + + return is_valid + + def _assert_keys_normalized( keys: set[ComboKey], input_shapes: tuple[tuple[int, ...], ...], @@ -1200,10 +1435,13 @@ def _compare_rules( variant: str, stats: ComparisonStats, sample: SampleInput | None = None, + untestable: set[ComboKey] | None = None, ) -> None: """Compare ground truth valid rules against DTensor claimed rules, updating stats.""" if not dtensor_rules: return + if untestable is None: + untestable = set() _assert_keys_normalized(ground_truth_valid, input_shapes, output_shapes) _assert_keys_normalized(dtensor_rules, input_shapes, output_shapes) @@ -1215,7 +1453,7 @@ def _compare_rules( stats.true_positives_by_op[op_str] = ( stats.true_positives_by_op.get(op_str, 0) + 1 ) - else: + elif combo_key not in untestable: stats.false_negatives.append( Discrepancy( input_placements=combo_key[0], @@ -1232,7 +1470,7 @@ def _compare_rules( ) for combo_key in dtensor_rules: - if combo_key not in ground_truth_valid: + if combo_key not in ground_truth_valid and combo_key not in untestable: stats.false_positives.append( Discrepancy( input_placements=combo_key[0], @@ -1374,12 +1612,201 @@ def _discover_aten_op( tensors = extract_tensors_from_sample(sample) if not tensors or any(0 in t.shape for _, t in tensors): continue - aten_op, _, _ = get_aten_op_for_sample(opinfo.op, sample, opinfo.name) + capture = get_aten_op_for_sample(opinfo.op, sample, opinfo.name) + aten_op = capture.best_match if aten_op is not None: return aten_op return None +def _check_ground_truth( + result: Any, +) -> torch.Tensor | list[torch.Tensor] | None: + """Validate an op result is suitable as ground truth. + + Returns the ground truth tensor(s) or None if the result should be skipped. + """ + if isinstance(result, (list, tuple)): + if not all(isinstance(t, torch.Tensor) for t in result): + return None + gt = list(result) + elif isinstance(result, torch.Tensor): + gt = result + else: + return None + + first_gt = gt[0] if isinstance(gt, list) else gt + if first_gt.numel() == 0: + return None + if (first_gt == 0).all(): + return None + if first_gt.isnan().all(): + return None + return gt + + +def _validate_aten_op_for_sample( + aten_op: OpOverload, + captured_args: tuple[Any, ...], + captured_kwargs: dict[str, Any], + ground_truth: torch.Tensor | list[torch.Tensor], + world_size: int, + incorrect_only: bool, + verbose: bool, + sample_idx: int, + variant: str, + stats: ComparisonStats, + sample: SampleInput | None = None, +) -> tuple[int, int]: + """Validate a single aten op with captured args against ground truth. + + Shared logic used by both default (1:1) and allow_composite modes in + compare_operator. Returns (samples_counted, combinations_counted). + """ + tensors = extract_tensors_from_args(captured_args, captured_kwargs) + if not tensors: + return 0, 0 + if any(0 in t.shape for _, t in tensors): + return 0, 0 + + input_shapes = tuple(t.shape for _, t in tensors) + gt_list = ground_truth if isinstance(ground_truth, list) else [ground_truth] + output_shapes = tuple(tuple(gt.shape) for gt in gt_list) + n_outputs = len(gt_list) + first_gt = gt_list[0] + + scalar_args = tuple(a for a in captured_args if not isinstance(a, torch.Tensor)) + scalar_kwargs = { + k: v for k, v in captured_kwargs.items() if not isinstance(v, torch.Tensor) + } + + mitigations = _prepare_aten_mitigations(aten_op, captured_args, captured_kwargs) + + input_placement_options = [ + get_1d_input_placements_for_tensor(t, include_partial=True) for _, t in tensors + ] + output_placement_options = get_1d_output_placements_for_tensor(first_gt) + + dtensor_rules = _query_dtensor_rules( + aten_op, + tensors, + captured_args, + captured_kwargs, + input_shapes, + output_shapes, + world_size, + verbose, + ) + + ground_truth_valid: set[ComboKey] = set() + total_combinations = 0 + + tensor_device = tensors[0][1].device.type if tensors else "cpu" + with LocalTensorMode(frozenset(range(world_size))): + mesh = init_device_mesh(tensor_device, (world_size,)) + + if incorrect_only: + combinations_to_test = [] + for combo_key in dtensor_rules: + input_plc_strs, output_plc_strs = combo_key + input_plcs_list: list[Placement] = [] + all_valid = True + for s in input_plc_strs: + p = parse_placement(s) + if p is None: + all_valid = False + break + input_plcs_list.append(p) + output_plcs_list: list[Placement] = [] + for s in output_plc_strs: + p = parse_placement(s) + if p is None: + all_valid = False + break + output_plcs_list.append(p) + if not all_valid: + continue + combinations_to_test.append( + ( + tuple(input_plcs_list), + tuple(output_plcs_list), + combo_key, + ) + ) + else: + combinations_to_test = [] + for input_placements in itertools.product(*input_placement_options): + if is_fully_replicated(input_placements): + continue + for output_placement in output_placement_options: + output_placements = tuple( + output_placement for _ in range(n_outputs) + ) + combo_key = ( + tuple(str(p) for p in input_placements), + tuple(str(p) for p in output_placements), + ) + combinations_to_test.append( + (input_placements, output_placements, combo_key) + ) + + untestable: set[ComboKey] = set() + + for ( + input_placements, + output_placements, + combo_key, + ) in combinations_to_test: + total_combinations += 1 + is_valid = _validate_aten_with_mitigations( + aten_op, + captured_args, + captured_kwargs, + input_placements, + output_placements, + ground_truth, + world_size, + mesh, + mitigations, + ) + + if is_valid is None: + normalized_key = normalize_combo_key( + combo_key, input_shapes, output_shapes + ) + untestable.add(normalized_key) + elif is_valid: + normalized_key = normalize_combo_key( + combo_key, input_shapes, output_shapes + ) + if not is_fully_replicated( + tuple(parse_placement(p) or Replicate() for p in normalized_key[0]) + ): + ground_truth_valid.add(normalized_key) + + _compare_rules( + ground_truth_valid, + dtensor_rules, + input_shapes, + output_shapes, + sample_idx, + scalar_args, + scalar_kwargs, + aten_op, + variant, + stats, + sample, + untestable, + ) + + if verbose: + print(f" Sample {sample_idx} [{aten_op}]: shapes={input_shapes}") + print(f" Ground truth valid: {len(ground_truth_valid)}") + print(f" DTensor rules: {len(dtensor_rules)}") + + return 1, total_combinations + + def compare_operator( op_name: str, device: str = "cpu", @@ -1388,6 +1815,7 @@ def compare_operator( max_samples: int | None = None, verbose: bool = False, incorrect_only: bool = False, + allow_composite: bool = False, ) -> ComparisonStats: """ Compare DTensor's sharding rules against ground truth for an operator. @@ -1400,7 +1828,10 @@ def compare_operator( max_samples: Maximum number of samples to test per OpInfo verbose: Print detailed output incorrect_only: If True, only test DTensor's claimed rules for correctness. - Skips exhaustive search for missing rules (much faster). + Skips search for missing rules (much faster). + allow_composite: If True, validate each supported aten op individually for + samples that decompose into multiple aten calls. Default (False) + skips samples where the OpInfo doesn't map 1:1 to a single aten op. """ if op_name in SKIP_OPS: return ComparisonStats() @@ -1409,10 +1840,15 @@ def compare_operator( stats = ComparisonStats() - aten_op = _discover_aten_op(opinfos, device, dtype) - if aten_op is None or not _has_dtensor_support(aten_op): - stats.no_dtensor_support = True - return stats + if not allow_composite: + aten_op = _discover_aten_op(opinfos, device, dtype) + if aten_op is None or not _has_dtensor_support(aten_op): + if verbose: + print(f" ATEN_OP_MAP: {op_name} -> {aten_op} [no_support]") + stats.no_dtensor_support = True + return stats + if verbose: + print(f" ATEN_OP_MAP: {op_name} -> {aten_op} [supported]") total_samples = 0 total_combinations = 0 @@ -1420,7 +1856,7 @@ def compare_operator( for opinfo in opinfos: variant = opinfo.variant_test_name - if variant: + if variant and verbose: print(f"\n OpInfo variant: {variant}") op = opinfo.op @@ -1428,205 +1864,99 @@ def compare_operator( try: samples = list(opinfo.sample_inputs(device, dtype)) except Exception as e: - print(f" Error generating samples: {e}") + if verbose: + print(f" Error generating samples: {e}") continue if max_samples: samples = samples[:max_samples] for sample_idx, sample in enumerate(samples): - tensors = extract_tensors_from_sample(sample) - - if len(tensors) == 0: + # Check that SampleInput has tensor inputs and no zero-sized tensors + sample_tensors = extract_tensors_from_sample(sample) + if len(sample_tensors) == 0: skip_reasons["no tensor inputs"] += 1 continue - - if any(0 in t.shape for _, t in tensors): + if any(0 in t.shape for _, t in sample_tensors): skip_reasons["zero-sized tensor"] += 1 continue - total_samples += 1 - - try: - ground_truth_raw = _run_op_on_sample(op, sample) - if isinstance(ground_truth_raw, (list, tuple)): - if not all(isinstance(t, torch.Tensor) for t in ground_truth_raw): - total_samples -= 1 - skip_reasons["non-tensor output"] += 1 - continue - ground_truth = list(ground_truth_raw) - elif isinstance(ground_truth_raw, torch.Tensor): - ground_truth = ground_truth_raw - else: - total_samples -= 1 - skip_reasons["non-tensor output"] += 1 - continue - - # For skip checks, use the first tensor (or the only tensor) - first_gt = ( - ground_truth[0] if isinstance(ground_truth, list) else ground_truth - ) - if first_gt.numel() == 0: - total_samples -= 1 - skip_reasons["zero-element output"] += 1 - continue - if (first_gt == 0).all(): - total_samples -= 1 - skip_reasons["all-zero output"] += 1 - continue - if first_gt.isnan().all(): - total_samples -= 1 - skip_reasons["all-NaN output"] += 1 - continue - except Exception: - skip_reasons["op raised exception"] += 1 + # Capture all aten ops dispatched for this sample + capture = get_aten_op_for_sample(op, sample, opinfo.name) + if capture.best_match is None: + skip_reasons["no aten op captured"] += 1 continue - input_shapes = tuple(t.shape for _, t in tensors) - gt_list = ground_truth if isinstance(ground_truth, list) else [ground_truth] - output_shapes = tuple(tuple(gt.shape) for gt in gt_list) - n_outputs = len(gt_list) - - scalar_args = tuple( - a for a in sample.args if not isinstance(a, torch.Tensor) - ) - scalar_kwargs = { - k: v - for k, v in sample.kwargs.items() - if not isinstance(v, torch.Tensor) - } - - mitigations = _prepare_false_positive_mitigations(op, sample, tensors) - - input_placement_options = [ - get_1d_input_placements_for_tensor(t, include_partial=True) - for _, t in tensors + # Count supported aten ops in the capture + supported_ops = [ + (func, args, kwargs, result) + for func, args, kwargs, result in capture.all_ops + if _has_dtensor_support(func) ] - # Use first output for enumerating placement options (DTensor applies - # the same placement to all outputs of multi-output ops) - output_placement_options = get_1d_output_placements_for_tensor(first_gt) - - aten_op, captured_args, captured_kwargs = get_aten_op_for_sample( - op, sample, opinfo.name - ) + num_supported = len(supported_ops) - dtensor_rules = _query_dtensor_rules( - aten_op, - tensors, - captured_args, - captured_kwargs, - input_shapes, - output_shapes, - world_size, - verbose, - ) + if allow_composite: + # Validate each supported aten op individually + if num_supported == 0: + skip_reasons["no supported aten ops"] += 1 + continue - ground_truth_valid: set[ComboKey] = set() - - tensor_device = tensors[0][1].device.type if tensors else "cpu" - with LocalTensorMode(frozenset(range(world_size))): - mesh = init_device_mesh(tensor_device, (world_size,)) - - if incorrect_only: - combinations_to_test = [] - for combo_key in dtensor_rules: - input_plc_strs, output_plc_strs = combo_key - input_plcs_list: list[Placement] = [] - all_valid = True - for s in input_plc_strs: - p = parse_placement(s) - if p is None: - all_valid = False - break - input_plcs_list.append(p) - output_plcs_list: list[Placement] = [] - for s in output_plc_strs: - p = parse_placement(s) - if p is None: - all_valid = False - break - output_plcs_list.append(p) - if not all_valid: - continue - combinations_to_test.append( - ( - tuple(input_plcs_list), - tuple(output_plcs_list), - combo_key, - ) - ) - else: - combinations_to_test = [] - for input_placements in itertools.product(*input_placement_options): - if is_fully_replicated(input_placements): - continue - for output_placement in output_placement_options: - # Apply same placement to all outputs (matches - # DTensor propagator behavior for multi-output ops) - output_placements = tuple( - output_placement for _ in range(n_outputs) - ) - combo_key = ( - tuple(str(p) for p in input_placements), - tuple(str(p) for p in output_placements), - ) - combinations_to_test.append( - (input_placements, output_placements, combo_key) - ) - - for ( - input_placements, - output_placements, - combo_key, - ) in combinations_to_test: - total_combinations += 1 - is_valid = _validate_with_mitigations( - op, - sample, - tensors, - input_placements, - output_placements, - ground_truth, + for func, args, kwargs, result in supported_ops: + gt = _check_ground_truth(result) + if gt is None: + skip_reasons["non-tensor/degenerate aten output"] += 1 + continue + n_samples, n_combos = _validate_aten_op_for_sample( + func, + args, + kwargs, + gt, world_size, - mesh, - mitigations, + incorrect_only, + verbose, + sample_idx, + variant, + stats, + sample, ) + total_samples += n_samples + total_combinations += n_combos + else: + # Default: only validate samples with a single supported aten op + if num_supported != 1: + skip_reasons["non-1:1 aten mapping"] += 1 + continue - if is_valid: - normalized_key = normalize_combo_key( - combo_key, input_shapes, output_shapes - ) - if not is_fully_replicated( - tuple( - parse_placement(p) or Replicate() - for p in normalized_key[0] - ) - ): - ground_truth_valid.add(normalized_key) - - _compare_rules( - ground_truth_valid, - dtensor_rules, - input_shapes, - output_shapes, - sample_idx, - scalar_args, - scalar_kwargs, - aten_op, - variant, - stats, - sample, - ) + func, args, kwargs, result = supported_ops[0] + gt = _check_ground_truth(result) + if gt is None: + skip_reasons["non-tensor/degenerate aten output"] += 1 + continue - if verbose: - print(f" Sample {sample_idx}: shapes={input_shapes}") - print(f" Ground truth valid: {len(ground_truth_valid)}") - print(f" DTensor rules: {len(dtensor_rules)}") + n_samples, n_combos = _validate_aten_op_for_sample( + func, + args, + kwargs, + gt, + world_size, + incorrect_only, + verbose, + sample_idx, + variant, + stats, + sample, + ) + total_samples += n_samples + total_combinations += n_combos stats.total_samples = total_samples stats.total_combinations = total_combinations stats.skip_reasons = dict(skip_reasons) + # In allow_composite mode, check DTensor support after processing + if allow_composite and total_samples == 0 and not skip_reasons: + stats.no_dtensor_support = True + return stats @@ -1772,6 +2102,12 @@ def _print_ops(label: str, ops: list) -> None: action="store_true", help="Only test DTensor's claimed rules (faster, skips missing detection)", ) + parser.add_argument( + "--allow-composite", + action="store_true", + help="Validate each supported aten op individually for decomposed ops " + "(default skips non-1:1 aten mappings)", + ) parser.add_argument("--device", default="cuda", help="Device to use") parser.add_argument("--dtype", default="float32", help="Dtype to use") parser.add_argument( @@ -1830,6 +2166,10 @@ def _print_ops(label: str, ops: list) -> None: # Preamble display_names = [_format_op_name(n) for n in op_names] print(f"Testing ops: {', '.join(display_names)}") + if args.allow_composite: + print( + "Mode: allow-composite (validates each aten op in decomposed samples)" + ) if args.incorrect_only: print("Mode: incorrect-only (fast)") print(f"Device: {args.device}, Dtype: {dtype}, World size: {args.world_size}") @@ -1854,7 +2194,9 @@ def _print_ops(label: str, ops: list) -> None: dtype, args.world_size, args.max_samples, + verbose=True, incorrect_only=args.incorrect_only, + allow_composite=args.allow_composite, ) elapsed = time.time() - op_start diff --git a/torch/distributed/tensor/_ops/utils.py b/torch/distributed/tensor/_ops/utils.py index 6bb6309feb1a6..1381a9abba6f8 100644 --- a/torch/distributed/tensor/_ops/utils.py +++ b/torch/distributed/tensor/_ops/utils.py @@ -4,7 +4,7 @@ import itertools import operator from collections.abc import Callable, Iterable, Sequence -from typing import cast, TypeAlias, TypeVar +from typing import TypeAlias, TypeVar import torch from torch._prims_common import DimsSequenceType, DimsType @@ -22,6 +22,7 @@ ) from torch.distributed.tensor.device_mesh import DeviceMesh from torch.distributed.tensor.placement_types import ( + _is_shard_like, _StridedShard, Partial, Placement, @@ -205,7 +206,7 @@ def is_tensor_shardable( # number of shards in each tensor dimension num_shards = [1] * len(shape) for i, placement in enumerate(spec.placements): - if isinstance(placement, Shard | _StridedShard): + if _is_shard_like(placement): shard_dim = placement.dim if shard_dim >= len(shape): return False @@ -229,7 +230,7 @@ def is_tensor_evenly_shardable(shape: Sequence[int], spec: DTensorSpec) -> bool: # number of shards in each tensor dimension num_shards = [1] * len(shape) for i, placement in enumerate(spec.placements): - if isinstance(placement, Shard | _StridedShard): + if _is_shard_like(placement): shard_dim = placement.dim if shard_dim >= len(shape): return False @@ -254,17 +255,24 @@ def is_tensor_evenly_shardable_on_dim( num_shards = 1 for i, placement in enumerate(spec.placements): - if placement.is_shard(): - shard_dim = cast(Shard, placement).dim - if shard_dim == dim: - num_shards *= spec.mesh.size(i) + if _is_shard_like(placement) and placement.dim == dim: + num_shards *= spec.mesh.size(i) + if isinstance(placement, _StridedShard): + # _StridedShard._split_tensor first chunks into split_factor + # groups, then into num_shards within each group, so the dim + # must be divisible by the product of both. This is stricter + # than the final num_shards check and implies it. Note: + # num_shards already includes spec.mesh.size(i) from this + # iteration, so the check covers the full shard count. + if shape[dim] % (placement.split_factor * num_shards) != 0: + return False return shape[dim] % num_shards == 0 def is_tensor_dim_sharded(spec: DTensorSpec, dim: int) -> bool: """Return True if tensor dim is sharded.""" - return any(p.is_shard(dim) for p in spec.placements) + return any(_is_shard_like(p) and p.dim == dim for p in spec.placements) def is_tensor_partial(spec: DTensorSpec) -> bool: @@ -309,7 +317,7 @@ def map_placements_after_broadcast( elif isinstance(placement, Replicate): new_placements.append(placement) else: - if not isinstance(placement, Shard | _StridedShard): + if not _is_shard_like(placement): raise AssertionError shard_dim = normalize_dim(placement.dim, len(shape)) new_shard_dim = broadcast_dims_map[shard_dim] @@ -367,6 +375,7 @@ def expand_to_full_mesh_op_strategy( [list[DTensorSpec], DTensorSpec | tuple[DTensorSpec | None, ...]], bool ] | None = None, + different_mesh_args: list[int] | None = None, ) -> OpStrategy: """ Convenience function to allow writing a sharding strategy considering only a single mesh dimension, @@ -403,6 +412,18 @@ def expand_to_full_mesh_op_strategy( args_strategy = op_schema.args_strategy kwargs_strategy = op_schema.kwargs_strategy input_args_strategy = args_strategy + kwargs_strategy + + # Propagate use_strided_shard_as_shard_order from inputs so that + # strategy specs with _StridedShard get the correct flag (and thus + # correct shard_order) at construction time, avoiding shard_order + # mismatches in redistribute_cost computation. + _input_use_strided: bool | None = None + for input_strat in input_args_strategy: + input_spec = input_strat.strategies[0].output_spec + if any(isinstance(p, _StridedShard) for p in input_spec.placements): + _input_use_strided = input_spec.use_strided_shard_as_shard_order + break + all_strategies = [] # Track input placements if we skip strategies due to inplace placement mismatch blocking_inplace_input_placements: tuple[Placement, ...] | None = None @@ -441,7 +462,20 @@ def expand_to_full_mesh_op_strategy( input_strategy_counter += 1 # pyrefly: ignore [bad-argument-type] - spec_list.append(DTensorSpec(mesh, specs, tensor_meta=tensor_meta)) + use_strided = ( + _input_use_strided + if _input_use_strided is not None + and any(isinstance(p, _StridedShard) for p in specs) + else None + ) + spec_list.append( + DTensorSpec( + mesh, + specs, + tensor_meta=tensor_meta, + use_strided_shard_as_shard_order=use_strided, + ) + ) else: spec_list.append(None) @@ -476,6 +510,52 @@ def expand_to_full_mesh_op_strategy( f"input_specs({len(input_specs)}) != strategies({len(input_args_strategy)}: " f"{len(args_strategy)} args + {len(kwargs_strategy)} kwargs)" ) + + # Note [Multi-mesh args] + # + # Some ops accept args whose DTensor lives on a different DeviceMesh + # than the op's primary compute mesh. We call these "multi-mesh + # args". They arise in fused optimizer ops (e.g. _fused_adam_) + # where *state_steps* is a per-rank scalar counter allocated on a + # smaller sub-mesh (e.g. 1-D DP) while params and grads live on a + # larger mesh (e.g. 2-D DP × TP). + # + # Why must these args be Replicate? + # Sharding implies a specific partitioning of a tensor's data + # across the ranks of a mesh. If a tensor doesn't even *exist* + # on the compute mesh, there is no meaningful way to interpret a + # Shard placement for it. Replicate, on the other hand, is + # mesh-agnostic: every rank already holds the full data, so the + # op can simply read the value regardless of which mesh owns it. + # + # What we do here: + # We preserve the original mesh and Replicate placement for these + # args so the propagator does not try to redistribute them onto + # the compute mesh (which would fail or produce wrong results). + # + # This is distinct from the *element_mesh* handling in + # single_dim_strategy.py, which deals with foreach ops where + # different *elements* in a tensor list may live on different + # sub-meshes (e.g. param group A on 2-D mesh, param group B on + # 1-D mesh). + # TODO: refactor fused_ops handling so that there are no longer + # args on different meshes + if different_mesh_args is not None: + for idx in different_mesh_args: + if idx < len(input_args_strategy): + cross_mesh_input = input_args_strategy[idx] + original_spec = cross_mesh_input.strategies[0].output_spec + if original_spec.mesh != mesh: + if not all(p == Replicate() for p in original_spec.placements): + raise RuntimeError( + f"Cross-mesh input at index {idx} must be Replicate, " + f"but got {original_spec.placements}" + ) + input_specs[idx] = DTensorSpec( + mesh=original_spec.mesh, + placements=original_spec.placements, + tensor_meta=original_spec.tensor_meta, + ) self_spec = input_args_strategy[0].strategies[0].output_spec redistribute_input = self_spec.placements != input_specs[0].placements @@ -564,7 +644,11 @@ def shift_shard_dims_after_insert( ) -> Sequence[Placement]: normalized_placements: list[Placement] = [] for placement in placements: - if isinstance(placement, Shard) and placement.dim >= insert_dim: + if isinstance(placement, _StridedShard) and placement.dim >= insert_dim: + normalized_placements.append( + _StridedShard(placement.dim + 1, split_factor=placement.split_factor) + ) + elif isinstance(placement, Shard) and placement.dim >= insert_dim: normalized_placements.append(Shard(placement.dim + 1)) else: normalized_placements.append(placement) @@ -576,7 +660,11 @@ def shift_shard_dims_after_remove( ) -> Sequence[Placement]: normalized_placements: list[Placement] = [] for placement in placements: - if isinstance(placement, Shard) and placement.dim > remove_dim: + if isinstance(placement, _StridedShard) and placement.dim > remove_dim: + normalized_placements.append( + _StridedShard(placement.dim - 1, split_factor=placement.split_factor) + ) + elif isinstance(placement, Shard) and placement.dim > remove_dim: normalized_placements.append(Shard(placement.dim - 1)) else: normalized_placements.append(placement) diff --git a/torch/distributed/tensor/_redistribute.py b/torch/distributed/tensor/_redistribute.py index 25417cbe0356e..3608e3edcbc2c 100644 --- a/torch/distributed/tensor/_redistribute.py +++ b/torch/distributed/tensor/_redistribute.py @@ -17,6 +17,7 @@ from torch.distributed._functional_collectives import _are_we_tracing from torch.distributed.tensor._collective_utils import one_step_redistribute_cost from torch.distributed.tensor._dtensor_spec import ( + _StridedShardNotDecodableError, DTensorSpec, ShardOrder, ShardOrderEntry, @@ -25,6 +26,7 @@ from torch.distributed.tensor._utils import assert_no_mixed_partial_types from torch.distributed.tensor.device_mesh import DeviceMesh from torch.distributed.tensor.placement_types import ( + _is_shard_like, _StridedShard, Partial, Placement, @@ -165,11 +167,11 @@ def _comm_type_key(self) -> str | None: src, dst = self.src_dst_placements if src.is_partial() and dst.is_replicate(): return "all_reduce" - elif src.is_partial() and dst.is_shard(): + elif src.is_partial() and _is_shard_like(dst): return "reduce_scatter" - elif src.is_shard() and dst.is_replicate(): + elif _is_shard_like(src) and dst.is_replicate(): return "all_gather" - elif src.is_shard() and dst.is_shard(): + elif _is_shard_like(src) and _is_shard_like(dst): return "all_to_all" else: # Local ops (Replicate->Shard, Replicate->Partial, noop, etc.) @@ -252,18 +254,14 @@ def _update_shard_order_and_placements( current_placements[mesh_dim] = dst_placement -def _get_flattened_mesh_by_layout( +def _get_flattened_mesh_by_layout_impl( mesh: DeviceMesh, mesh_dims: tuple[int, ...] ) -> DeviceMesh | None: """ Query for an explicitly created flattened mesh using layout comparison. - Args: - mesh: The DeviceMesh to query - mesh_dims: Tuple of mesh dimension indices to look for - - Returns: - The flattened DeviceMesh if it was explicitly created, None otherwise. + Searches root_mesh._flatten_mapping for a mesh whose layout matches + the expected flattened layout for the given dims. Pure Python layout math. """ root_mesh = mesh._get_root_mesh() mesh_dim_names = mesh.mesh_dim_names @@ -289,6 +287,29 @@ def _get_flattened_mesh_by_layout( return None +def _get_flattened_mesh_by_layout( + mesh: DeviceMesh, mesh_dims: tuple[int, ...] +) -> DeviceMesh | None: + """ + Query for an explicitly created flattened mesh using layout comparison. + + When tracing with compile_on_one_rank, delegates to a custom op so the + flattened mesh appears as a call_function node derived from mesh (a graph + input) rather than as a get_attr constant holding an unpicklable + ProcessGroup. + """ + if _are_we_tracing() and torch.distributed.config.compile_on_one_rank: + # Pre-check: the custom op can't return None (torch.library doesn't + # support Optional opaque return types), so guard here first. + if _get_flattened_mesh_by_layout_impl(mesh, mesh_dims) is None: + return None + from torch.distributed._ops import device_mesh as _ # noqa: F401 + + return torch.ops.device_mesh._get_flattened_submesh(mesh, list(mesh_dims)) + + return _get_flattened_mesh_by_layout_impl(mesh, mesh_dims) + + # Track (mesh_hash, mesh_dims, reason) we've already warned about to avoid repeated warnings _warned_flatten_issues: set[tuple[int, tuple[int, ...], str]] = set() @@ -1352,6 +1373,11 @@ def generate_greedy_transform_infos( target = target_placements[mesh_dim] # If target is not Shard, we can directly redistribute since we # are traversing from inner to outer placements here + # TODO: extend nested sharding detection to _StridedShard + # (isinstance check and is_shard() below miss it). + # Safe today: strategies convert _StridedShard to Replicate + # on ALL mesh dims for a given reduction dim, so misaligned + # nested _StridedShard targets can't arise. if isinstance(target, Shard): # If target is Shard, check for nested sharding on the # tensor dim BEFORE the current mesh_dim @@ -1438,9 +1464,22 @@ def _gen_transform_infos_non_cached( src_spec.tensor_meta, ) if use_graph_based_transform: - transform_infos = drp.generate_graph_based_transform_infos( - src_spec, dst_spec, src_spec.shape - ) + # TODO(zpcore): Temporary workaround for the case where _StridedShard + # cannot be decoded into shard order. This happens when + # use_strided_shard_as_shard_order defaults to True (e.g. in + # Redistribute.forward where the target DTensorSpec is constructed from + # raw placements without the flag), but the split_factor doesn't + # correspond to any valid product of mesh dimension sizes (e.g. sf=2 + # on a 1D mesh). A proper fix is to either pass + # use_strided_shard_as_shard_order through the Redistribute API, or + # migrate to explicit shard_order so _StridedShard is no longer + # overloaded for two purposes. + try: + transform_infos = drp.generate_graph_based_transform_infos( + src_spec, dst_spec, src_spec.shape + ) + except _StridedShardNotDecodableError: + transform_infos = drp.generate_greedy_transform_infos(src_spec, dst_spec) else: transform_infos = drp.generate_greedy_transform_infos(src_spec, dst_spec) return transform_infos @@ -1542,6 +1581,15 @@ def redistribute_local_tensor( mesh_to_use = device_mesh i = transform_info.mesh_dim current, target = transform_info.src_dst_placements + + # _StridedShard methods use device_mesh directly, not mesh_to_use. + # This is safe because _StridedShard.is_shard() returns False, so + # _comm_type_key() returns None and flattening is never attempted. + if isinstance(current, _StridedShard) or isinstance(target, _StridedShard): + assert mesh_to_use is device_mesh, ( # noqa: S101 + "_StridedShard redistribute assumes no flattened transforms" + ) + num_chunks = mesh_to_use.size(mesh_dim=i) if current == target: @@ -1612,8 +1660,15 @@ def redistribute_local_tensor( target_placement.dim, ) elif isinstance(current, _StridedShard): - raise NotImplementedError( - "Redistribute from _StridedShard to Shard is not implemented yet" + # _StridedShard -> Shard: go via Replicate as intermediate + replicated = current._to_replicate_tensor( + local_tensor, device_mesh, i, transform_info.logical_shape + ) + new_local_tensor = target_placement._replicate_to_shard( + replicated, + mesh_to_use, + i, + mesh_to_use._sym_get_coordinate(i), ) else: raise ValueError( @@ -1625,7 +1680,7 @@ def redistribute_local_tensor( new_local_tensor = partial_spec._partition_value( local_tensor, mesh_to_use, i ) - elif current.is_shard() or isinstance(current, _StridedShard): + elif _is_shard_like(current): raise RuntimeError( f"redistribute from {current} to {target} not supported yet" ) @@ -1639,8 +1694,13 @@ def redistribute_local_tensor( elif isinstance(target, _StridedShard): # Case 4: target is _StridedShard if current.is_partial(): - raise NotImplementedError( - "Redistribute from Partial to _StridedShard is not implemented yet" + # Partial -> _StridedShard: reduce to Replicate, then strided shard + partial_spec = cast(Partial, current) + replicated = partial_spec._reduce_value( + local_tensor, mesh_to_use, i + ) + new_local_tensor = target._replicate_to_strided_shard( + replicated, device_mesh, i, device_mesh._sym_get_coordinate(i) ) elif current.is_replicate(): # split the tensor and return the corresponding local strided shard @@ -1648,9 +1708,13 @@ def redistribute_local_tensor( local_tensor, device_mesh, i, device_mesh._sym_get_coordinate(i) ) elif current.is_shard(): - # Shard -> _StridedShard on potentially different dimensions - raise NotImplementedError( - "Redistribute from Shard to _StridedShard is not implemented yet" + # Shard -> _StridedShard: all-gather to Replicate, then strided shard + current_placement = cast(Shard, current) + replicated = current_placement._to_replicate_tensor( + local_tensor, mesh_to_use, i, transform_info.logical_shape + ) + new_local_tensor = target._replicate_to_strided_shard( + replicated, device_mesh, i, device_mesh._sym_get_coordinate(i) ) elif isinstance(current, _StridedShard): # _StridedShard -> _StridedShard: go through Replicate @@ -1674,6 +1738,99 @@ def redistribute_local_tensor( return new_local_tensor +def _redistribute_backward( + grad_output: "dtensor.DTensor", + previous_spec: DTensorSpec, + original_dtype: torch.dtype | None = None, + backward_dtype: torch.dtype | None = None, + async_op: bool = False, +): + """ + Common function for redistributing a distributed tensor during backward + and twice-backward backpropagation steps. + + Args: + grad_output: The output gradient tensor. + previous_spec: DTensorSpec prior to redistribution. + original_dtype: Original output tensor dtype from forward pass (for type checking) + backward_dtype: Desired data type for backwards output. + async_op: whether to perform the DTensor redistribute operation + asynchronously or not. Default: False + + Returns: + A :class:`torch.Tensor` object. + A :class:`DTensorSpec` object. + """ + if backward_dtype is not None and backward_dtype != grad_output._local_tensor.dtype: + local_tensor = grad_output._local_tensor.to(dtype=backward_dtype) + current_spec = DTensorSpec( + mesh=grad_output._spec.device_mesh, + placements=grad_output._spec.placements, + tensor_meta=TensorMeta( + shape=grad_output.shape, + stride=grad_output.stride(), + # pyrefly: ignore [bad-argument-type] + dtype=backward_dtype, + ), + use_strided_shard_as_shard_order=grad_output._spec.use_strided_shard_as_shard_order, + ) + previous_spec = DTensorSpec( + mesh=previous_spec.device_mesh, + placements=previous_spec.placements, + tensor_meta=current_spec.tensor_meta, + use_strided_shard_as_shard_order=previous_spec.use_strided_shard_as_shard_order, + ) + else: + local_tensor = grad_output._local_tensor + current_spec = grad_output._spec + # skip the replicate to partial transformation when we are in backward pass + # In this case we keep the grad as replicate, this is because we don't + # want to convert the replicated gradients back to partial, although + # that's logically conform with the same layout, converting the gradients + # back to partial is actually useless as you would have to do reduce later + # which would be more expensive than keeping it replicate! + + # for backward shard -> partial, we just do shard -> replicate + # for backward replicate -> partial, we skip the transformation + # NOTE: _is_shard_like covers _StridedShard defensively; currently + # unreachable because Partial -> _StridedShard is not implemented. + normalized_placements: list[Placement] = [] + for current, target in zip(current_spec.placements, previous_spec.placements): + if (_is_shard_like(current) or current.is_replicate()) and target.is_partial(): + normalized_placements.append(Replicate()) + else: + normalized_placements.append(target) + + previous_spec = DTensorSpec( + previous_spec.device_mesh, + placements=tuple(normalized_placements), + tensor_meta=previous_spec.tensor_meta, + use_strided_shard_as_shard_order=previous_spec.use_strided_shard_as_shard_order, + ) + + output = redistribute_local_tensor( + local_tensor, + current_spec, + previous_spec, + async_op=async_op, + ) + + if output.dtype != original_dtype: + output = output.to(original_dtype) + + spec = DTensorSpec( + previous_spec.device_mesh, + tuple(normalized_placements), + tensor_meta=TensorMeta( + shape=grad_output.shape, + stride=grad_output.stride(), + dtype=output.dtype, + ), + use_strided_shard_as_shard_order=previous_spec.use_strided_shard_as_shard_order, + ) + return output, spec + + class Redistribute(torch.autograd.Function): @staticmethod def forward( # type: ignore[override] @@ -1700,6 +1857,7 @@ def forward( # type: ignore[override] stride=input.stride(), dtype=forward_dtype, ), + use_strided_shard_as_shard_order=input._spec.use_strided_shard_as_shard_order, ) else: local_tensor = input._local_tensor @@ -1736,72 +1894,60 @@ def forward( # type: ignore[override] @staticmethod def backward(ctx, grad_output: "dtensor.DTensor"): # type: ignore[override] previous_spec = ctx.current_spec - async_op = ctx.async_op - backward_dtype = ctx.backward_dtype or ctx.original_dtype + output_dtensor = NestedRedistribute.apply( + grad_output, + previous_spec, + ctx.async_op, + ctx.backward_dtype, + ctx.original_dtype, + ) + return ( + output_dtensor, + None, + None, + None, + None, + None, + ) - if backward_dtype != grad_output._local_tensor.dtype: - local_tensor = grad_output._local_tensor.to(dtype=backward_dtype) - current_spec = DTensorSpec( - mesh=grad_output._spec.device_mesh, - placements=grad_output._spec.placements, - tensor_meta=TensorMeta( - shape=grad_output.shape, - stride=grad_output.stride(), - dtype=backward_dtype, - ), - ) - previous_spec = DTensorSpec( - mesh=previous_spec.device_mesh, - placements=previous_spec.placements, - tensor_meta=current_spec.tensor_meta, - ) - else: - local_tensor = grad_output._local_tensor - current_spec = grad_output._spec - # skip the replicate to partial transformation when we are in backward pass - # In this case we keep the grad as replicate, this is because we don't - # want to convert the replicated gradients back to partial, although - # that's logically conform with the same layout, converting the gradients - # back to partial is actually useless as you would have to do reduce later - # which would be more expensive than keeping it replicate! - - # for backward shard -> partial, we just do shard -> replicate - # for backward replicate -> partial, we skip the transformation - normalized_placements: list[Placement] = [] - for current, target in zip(current_spec.placements, previous_spec.placements): - if (current.is_shard() or current.is_replicate()) and target.is_partial(): - normalized_placements.append(Replicate()) - else: - normalized_placements.append(target) - previous_spec = DTensorSpec( - previous_spec.device_mesh, - placements=tuple(normalized_placements), - tensor_meta=previous_spec.tensor_meta, - ) +class NestedRedistribute(torch.autograd.Function): + """ + This class is used to make the redistribution of a DTensor twice-differentiable. + This is called during the `Redistribute.forward`. + Therefore, `NestedRedistribute.forward` is called during the first backward pass, + and `NestedRedistribute.backward` is called during the second backward pass. - output = redistribute_local_tensor( - local_tensor, - current_spec, + Note: `NestedRedistribute.backward` is not differentiable, and therefore triple + backward is not yet supported. + """ + + @staticmethod + def forward( # type: ignore[override] + # pyre-fixme[2]: Parameter must be annotated. + ctx, + grad_output: "dtensor.DTensor", + previous_spec: DTensorSpec, + async_op: bool = False, + forward_dtype: torch.dtype | None = None, + backward_dtype: torch.dtype | None = None, + ): + ctx.async_op = async_op + ctx.original_dtype = grad_output._local_tensor.dtype + ctx.backward_dtype = backward_dtype or ctx.original_dtype + + output, spec = _redistribute_backward( + grad_output, previous_spec, - async_op=async_op, - is_explicit=True, + ctx.backward_dtype, + backward_dtype, + async_op, ) - if output.dtype != ctx.original_dtype: - output = output.to(ctx.original_dtype) + ctx.current_spec = spec - spec = DTensorSpec( - previous_spec.device_mesh, - tuple(normalized_placements), - tensor_meta=TensorMeta( - shape=grad_output.shape, - stride=grad_output.stride(), - dtype=output.dtype, - ), - ) # pyrefly: ignore [bad-argument-type] - output_dtensor = dtensor.DTensor( + return dtensor.DTensor( # pyrefly: ignore [bad-argument-count] output, spec, @@ -1809,6 +1955,20 @@ def backward(ctx, grad_output: "dtensor.DTensor"): # type: ignore[override] requires_grad=grad_output.requires_grad, ) + @staticmethod + def backward(ctx, grad2_output: "dtensor.DTensor"): # type: ignore[override] + previous_spec = ctx.current_spec + async_op = ctx.async_op + backward_dtype = ctx.backward_dtype or ctx.original_dtype + + output_dtensor = NestedRedistribute.apply( + grad2_output, + previous_spec, + async_op, + backward_dtype, + ctx.original_dtype, + ) + return ( output_dtensor, None, diff --git a/torch/distributed/tensor/_sharding_prop.py b/torch/distributed/tensor/_sharding_prop.py index a85b19130f69f..463a0d4362fc0 100644 --- a/torch/distributed/tensor/_sharding_prop.py +++ b/torch/distributed/tensor/_sharding_prop.py @@ -45,6 +45,56 @@ log = logging.getLogger(__name__) +def _propagate_use_strided_shard_flag( + op_strategy: OpStrategy, + op_schema: OpSchema, +) -> None: + """Propagate use_strided_shard_as_shard_order from input specs to output specs. + + When inputs carry _StridedShard with an explicit flag, all output (and input) + DTensorSpecs in the strategy that also contain _StridedShard must agree. + Strategy functions may forget to propagate the flag; this function fixes + them up centrally after the strategy is produced. + """ + _use_strided: bool | None = None + for spec in op_schema.args_spec: + if any(isinstance(p, _StridedShard) for p in spec.placements): + val = spec.use_strided_shard_as_shard_order + if _use_strided is not None and _use_strided != val: + raise ValueError( + "Conflicting use_strided_shard_as_shard_order across " + f"input specs: got both {_use_strided} and {val}" + ) + _use_strided = val + + if _use_strided is None: + return + + def _fixup(spec: DTensorSpec) -> None: + if not any(isinstance(p, _StridedShard) for p in spec.placements): + return + if spec.use_strided_shard_as_shard_order == _use_strided: + return + spec.use_strided_shard_as_shard_order = _use_strided + if _use_strided: + spec.shard_order = None # pyrefly: ignore[bad-assignment] + else: + spec.shard_order = DTensorSpec.compute_default_shard_order(spec.placements) + + for op_spec in op_strategy.strategies: + out = op_spec.output_specs + if out is not None: + if isinstance(out, DTensorSpec): + _fixup(out) + else: + for s in out: + if s is not None: + _fixup(s) + if op_spec.input_specs is not None: + for s in op_spec.input_specs: + _fixup(s) + + def _length(obj) -> int: if obj is None: return 0 @@ -53,14 +103,15 @@ def _length(obj) -> int: return len(obj) -def _get_expected_num_tensor_outputs(op: OpOverload) -> int: +def _get_expected_num_tensor_outputs(op: OpOverload) -> int | None: """ Get the expected number of tensor outputs for an operator based on its schema. Returns: The number of tensor outputs expected. Returns 0 for ops that don't return tensors (e.g., _linalg_check_errors). Returns 1 for single tensor return, and >1 for - tuple returns where each element is a tensor. + tuple returns where each element is a tensor. Returns None for List[Tensor] + returns where the length is unknown at schema time. """ return_types = op._schema.returns if len(return_types) == 0: @@ -71,8 +122,8 @@ def _get_expected_num_tensor_outputs(op: OpOverload) -> int: # Could be single tensor or tuple of tensors return len(return_types) elif isinstance(first_return.type, torch.ListType): - # List[Tensor] - we don't know the length at schema time, treat as 1 - return 1 + # List[Tensor] - we don't know the length at schema time + return None else: # Not a tensor return type return 0 @@ -101,6 +152,16 @@ def _validate_tensor_meta_count( else: actual_outputs = len(tensor_meta) + if expected_outputs is None: + # List[Tensor] return type: length unknown at schema time, but + # tensor_meta must be a list of TensorMeta. + if not isinstance(tensor_meta, list): + raise AssertionError( + f"Tensor meta for {op_schema.op} should be a list[TensorMeta] " + f"(op returns List[Tensor]), but got {type(tensor_meta).__name__}" + ) + return + if actual_outputs != expected_outputs: raise AssertionError( f"Tensor meta count mismatch for {op_schema.op}: " @@ -346,6 +407,15 @@ def __init__(self) -> None: aten.select_backward.default: 1, aten.slice_backward.default: 1, } + # squeeze ops that need dim arg rewritten to only globally-singleton dims + self.squeeze_op_to_dims_variant: dict[OpOverload, OpOverload] = { + aten.squeeze.default: aten.squeeze.dims, + aten.squeeze.dim: aten.squeeze.dims, + aten.squeeze.dims: aten.squeeze.dims, + aten.squeeze_.default: aten.squeeze_.dims, + aten.squeeze_.dim: aten.squeeze_.dims, + aten.squeeze_.dims: aten.squeeze_.dims, + } def register_sharding_prop_rule( self, @@ -541,21 +611,11 @@ def _create_output_spec_with_new_tensor_meta( if isinstance(spec, DTensorSpec): output_tensor_meta_i = output_tensor_meta[i] if not isinstance(output_tensor_meta_i, TensorMeta): - # NOTE: aten.convolution_backward.default is an exception and it - # needs extra handling because any Tensor in the output tuple - # can be `None` depending on the output_mask parameter. This can - # occur during double backpropagation or when certain gradients - # are not needed (e.g., grad_input when input has requires_grad=False, - # grad_weight/grad_bias when weight/bias have requires_grad=False, - # or grad_bias when bias is None). We explicitly allow the - # corresponding TensorMeta to be `None`. - if ( - op == aten.convolution_backward.default - and i in (0, 1, 2) - and output_tensor_meta_i is None - ): - if not isinstance(output_specs, list): - raise AssertionError + # Some ops (e.g. convolution_backward, native_layer_norm_backward, + # _fused_rms_norm_backward) have an output_mask parameter that + # controls which outputs are computed. When output_mask[i] is + # False, the output at position i is None and has no TensorMeta. + if output_tensor_meta_i is None: new_specs.append(None) continue else: @@ -713,6 +773,7 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin if op_strategy is not None: if isinstance(op_strategy, OpStrategy): + _propagate_use_strided_shard_flag(op_strategy, op_schema) # single Op strategy output_strategy = _select_min_cost_strategy(op_strategy, op_schema) @@ -770,6 +831,15 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin needs_redistribute = True use_val_from_redistribute_schema = True + # rewrite squeeze to use only globally-singleton dims + if op_schema.op in self.squeeze_op_to_dims_variant: + schema = suggestion_schema or op_schema + adjusted = self._adjust_squeeze_to_global_singletons(schema) + if adjusted is not None: + suggestion_schema = adjusted + needs_redistribute = True + use_val_from_redistribute_schema = True + # construct output spec for the op if op_schema.return_type_tuple_tensor_like(): # for ops that return multiple tensors and the output_specs is not @@ -784,6 +854,7 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin mesh=output_specs.mesh, placements=output_specs.placements, tensor_meta=output_specs.tensor_meta, + use_strided_shard_as_shard_order=output_specs.use_strided_shard_as_shard_order, ) for _ in range(len(op_schema.op._schema.returns)) ) @@ -809,9 +880,11 @@ def propagate_op_sharding_non_cached(self, op_schema: OpSchema) -> OutputShardin for strategy in op_strategy.children: if not isinstance(strategy, OpStrategy): raise AssertionError + _propagate_use_strided_shard_flag(strategy, op_schema) selected_strategy = _select_min_cost_strategy(strategy) selected_strategies.append(selected_strategy) - out_spec_list.append(selected_strategy.output_spec) + if selected_strategy.output_specs is not None: + out_spec_list.append(selected_strategy.output_spec) needs_redistribute = False suggestion_args: list[object] = [] @@ -957,3 +1030,47 @@ def _adjust_shape_and_stride_args( ) return OpSchema(schema.op, tuple(expected_input_schema), schema.kwargs_schema) + + def _adjust_squeeze_to_global_singletons(self, schema: OpSchema) -> OpSchema | None: + """ + Rewrite squeeze ops to squeeze.dims with only globally-singleton dims. + Fixes bug where sharded dims with local size 1 get incorrectly squeezed. + Returns None if no rewrite is needed (already squeeze.dims with correct args). + """ + from torch.fx.experimental.symbolic_shapes import guard_or_false + + input_spec = cast(DTensorSpec, schema.args_schema[0]) + tensor_meta = input_spec.tensor_meta + if tensor_meta is None: + raise RuntimeError("squeeze requires tensor metadata") + global_shape = tensor_meta.shape + ndim = len(global_shape) + + def normalize(d: int) -> int: + return d if d >= 0 else d + ndim + + def is_singleton(d: int) -> bool: + nd = normalize(d) + return 0 <= nd < ndim and guard_or_false(global_shape[nd] == 1) + + # guard_or_false: conservatively keep dims when size is symbolic/unknown + if schema.op in (aten.squeeze.default, aten.squeeze_.default): + target_dims = tuple( + i for i, s in enumerate(global_shape) if guard_or_false(s == 1) + ) + elif schema.op in (aten.squeeze.dim, aten.squeeze_.dim): + dim = normalize(schema.args_schema[1]) # type: ignore[arg-type] + target_dims = (dim,) if is_singleton(dim) else () + else: + dims = cast(Sequence[int], schema.args_schema[1]) + target_dims = tuple( # type: ignore[union-attr] + normalize(d) for d in dims if is_singleton(d) + ) + + dims_variant = self.squeeze_op_to_dims_variant[schema.op] + # Skip rewrite if already targeting the right op with the same dims + if schema.op == dims_variant and len(schema.args_schema) > 1: + existing_dims = schema.args_schema[1] + if existing_dims == target_dims: + return None + return OpSchema(dims_variant, (input_spec, target_dims), {}) diff --git a/torch/distributed/tensor/_utils.py b/torch/distributed/tensor/_utils.py index e326cd827974e..89aede42e53f5 100644 --- a/torch/distributed/tensor/_utils.py +++ b/torch/distributed/tensor/_utils.py @@ -1,7 +1,7 @@ import logging import threading from collections.abc import Callable, Sequence -from typing import Any, cast +from typing import Any import torch import torch.distributed._functional_collectives as funcol @@ -15,6 +15,7 @@ from torch.distributed.tensor._dtensor_spec import DTensorSpec from torch.distributed.tensor._op_schema import OpSchema from torch.distributed.tensor.placement_types import ( + _is_shard_like, _StridedShard, Partial, Placement, @@ -314,14 +315,13 @@ def compute_local_tensor_info( for idx, placement in enumerate(placements): mesh_dim_size = mesh.size(idx) - if placement.is_shard(): - shard_placement = cast(Shard, placement) - if shard_placement.dim < 0: + if _is_shard_like(placement): + if placement.dim < 0: raise AssertionError( "Shard placements should have negative dims normalized in " - f"the user-facing APIs: {shard_placement}" + f"the user-facing APIs: {placement}" ) - shard_dim = shard_placement.dim + shard_dim = placement.dim if shard_dim >= len(local_shape): raise AssertionError( f"Sharding dim {shard_dim} greater than tensor ndim {len(local_shape)} " @@ -385,6 +385,7 @@ def compute_global_tensor_shape( if isinstance(placements[0], Replicate): return shape + # NOTE: isinstance(_, Shard) does not match _StridedShard; see _is_shard_like(). elif isinstance(placements[0], Shard): @maybe_run_for_local_tensor diff --git a/torch/distributed/tensor/experimental/_context_parallel/_load_balancer.py b/torch/distributed/tensor/experimental/_context_parallel/_load_balancer.py index fb01a8c20d1da..f5d982d5a6561 100644 --- a/torch/distributed/tensor/experimental/_context_parallel/_load_balancer.py +++ b/torch/distributed/tensor/experimental/_context_parallel/_load_balancer.py @@ -148,27 +148,16 @@ def _generate_indices(self, restore: bool = False) -> Tensor: if seq_length % (world_size * 2) != 0: raise AssertionError chunk_size = seq_length // (world_size * 2) - all_indices = [] - for rank in range(world_size): - # Generate indices for first chunk of the cp rank - first_chunk_start = rank * chunk_size - first_chunk_indices = list( - range(first_chunk_start, first_chunk_start + chunk_size) - ) - - # Second chunk: positions from the complementary chunk - second_chunk_idx = world_size * 2 - rank - 1 - second_chunk_start = second_chunk_idx * chunk_size - second_chunk_indices = list( - range(second_chunk_start, second_chunk_start + chunk_size) - ) - # combine the indices for this rank - all_indices.extend(first_chunk_indices + second_chunk_indices) + # Split sequence into 2*world_size chunks, then pair chunk r with + # chunk (2*world_size - 1 - r) for each rank. + indices = torch.arange(seq_length, dtype=torch.int, device=self.device) + chunks = indices.view(world_size * 2, chunk_size) + head_idx = torch.arange(world_size, device=self.device) + tail_idx = 2 * world_size - 1 - head_idx + paired = torch.stack([chunks[head_idx], chunks[tail_idx]], dim=1) + all_indices_tensor = paired.reshape(-1) - all_indices_tensor = torch.tensor( - all_indices, dtype=torch.int, device=self.device - ) if restore: all_indices_tensor = torch.argsort(all_indices_tensor) diff --git a/torch/distributed/tensor/experimental/_tp_transform.py b/torch/distributed/tensor/experimental/_tp_transform.py index e78222845339c..adde891c4b68c 100644 --- a/torch/distributed/tensor/experimental/_tp_transform.py +++ b/torch/distributed/tensor/experimental/_tp_transform.py @@ -430,6 +430,7 @@ def _partition_val(val: Any, spec: DTensorSpec) -> Any: return local_shard for idx, placement in enumerate(spec.placements): + # NOTE: is_shard() does not match _StridedShard; see _is_shard_like(). if placement.is_shard(): placement = cast(Shard, placement) num_chunks = spec.mesh.size(mesh_dim=idx) diff --git a/torch/distributed/tensor/parallel/fsdp.py b/torch/distributed/tensor/parallel/fsdp.py index 99e7113afefb6..5f6facf74542a 100644 --- a/torch/distributed/tensor/parallel/fsdp.py +++ b/torch/distributed/tensor/parallel/fsdp.py @@ -37,6 +37,7 @@ def _get_box(tensor: DTensor) -> tuple[torch.Size, torch.Size]: offsets = [0] * len(tensor.size()) num_chunks = device_mesh.size(mesh_dim=0) + # NOTE: is_shard() does not match _StridedShard; see _is_shard_like(). if tensor.placements[0].is_shard(): shard_dim = cast(DShard, placement).dim chunk_size = tensor.size(shard_dim) // num_chunks @@ -81,6 +82,7 @@ def _create_sharded_tensor_md_from_dt( my_rank = dist.get_rank(dt_pg) scapegoat_rank = 0 if my_rank > 0 else 1 + # NOTE: is_shard() does not match _StridedShard; see _is_shard_like(). if dt.placements[0].is_shard(): shard_count = dt_pg.size() else: diff --git a/torch/distributed/tensor/parallel/input_reshard.py b/torch/distributed/tensor/parallel/input_reshard.py index 81e25621e040a..5e60cfcc7fb8f 100644 --- a/torch/distributed/tensor/parallel/input_reshard.py +++ b/torch/distributed/tensor/parallel/input_reshard.py @@ -4,6 +4,7 @@ import torch from torch.distributed.tensor import DeviceMesh, DTensor, Replicate, Shard +from torch.distributed.tensor.placement_types import _is_shard_like __all__ = [ @@ -87,7 +88,7 @@ def _unpack_hook_tp(mesh: DeviceMesh, input_reshard_dim: int, x: Any) -> torch.T if ( isinstance(x, DTensor) and len(x._spec.placements) == 1 - and x._spec.placements[0].is_shard() + and _is_shard_like(x._spec.placements[0]) ): return x.redistribute(device_mesh=mesh, placements=[Replicate()]) elif ( diff --git a/torch/distributed/tensor/placement_types.py b/torch/distributed/tensor/placement_types.py index 53b96e9051dd4..67ab7b3f12754 100644 --- a/torch/distributed/tensor/placement_types.py +++ b/torch/distributed/tensor/placement_types.py @@ -4,7 +4,7 @@ import functools from collections.abc import Sequence from dataclasses import dataclass, field -from typing import cast, TypeVar +from typing import cast, TypeGuard, TypeVar import torch import torch._C @@ -62,6 +62,10 @@ class Shard(torch._C._distributed.Shard): .. warning:: sharding on a tensor dimension where the tensor dimension size is not evenly divisible on a DeviceMesh dimension is currently experimental and subject to change. + + .. note:: When checking whether a placement is shard-like, use + :func:`_is_shard_like` instead of ``isinstance(p, Shard)`` to also + match :class:`_StridedShard`. """ def _split_tensor( @@ -255,7 +259,7 @@ def _local_shard_size_and_offset( num_chunks: int, rank: RankType, ) -> tuple[int, RankType]: - # pyrefly: ignore[bad-argument-type] # pyrefly bug + # pyrefly: ignore [bad-argument-type, bad-return] return Shard.local_shard_size_and_offset(curr_local_size, num_chunks, rank) @staticmethod @@ -419,7 +423,7 @@ def _maybe_unpad_tensor( logical_dim_size: IntLikeType, num_chunks: int, ) -> torch.Tensor: - from torch.fx.experimental.symbolic_shapes import guard_or_true + from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true # Assume padding (uneven sharding) as general case for unbacked sizes. is_padded = guard_or_true(logical_dim_size % num_chunks != 0) @@ -429,6 +433,22 @@ def _maybe_unpad_tensor( unpad_size = full_chunk_size * num_chunks - logical_dim_size # type: ignore[possibly-undefined] local_tensor = unpad_tensor(local_tensor, self.dim, unpad_size) + # Bind derived symbolic sizes (e.g. 2*(s//2)) back to the original + # symbol - needed for correct shape propagation and dynamo generation + if local_tensor.size(self.dim) is not logical_dim_size: + orig_size = local_tensor.size(self.dim) + torch._check(orig_size >= logical_dim_size) + local_tensor = local_tensor.narrow(self.dim, 0, logical_dim_size) + + # Safety check: the narrow should never change the concrete size. + # Use guard_or_false so we don't trigger data-dependent guards + # on unbacked symints. + if guard_or_false(local_tensor.size(self.dim) != orig_size): + raise RuntimeError( + f"narrow unexpectedly changed concrete size on dim {self.dim}: " + f"{orig_size} -> {local_tensor.size(self.dim)}" + ) + return local_tensor def _to_replicate_tensor( @@ -1158,6 +1178,17 @@ def local_shard_size_and_offset( return local_shard_size, offsets +def _is_shard_like(p: "Placement") -> TypeGuard[Shard | _StridedShard]: + """Check if a placement is Shard or _StridedShard. + + Use this instead of ``isinstance(p, Shard)`` to avoid silently missing + ``_StridedShard``. When ``_StridedShard`` is unified with ``Shard`` + (see TODO on the class), this helper can be collapsed to a single + ``isinstance`` check. + """ + return isinstance(p, Shard | _StridedShard) + + class Replicate(torch._C._distributed.Replicate): """ The ``Replicate()`` placement describes the DTensor replicating on a corresponding diff --git a/torch/distributed/utils.py b/torch/distributed/utils.py index 28a18428c673d..d5651926b102a 100644 --- a/torch/distributed/utils.py +++ b/torch/distributed/utils.py @@ -134,16 +134,16 @@ def to_map(obj): from torch.nn.parallel.scatter_gather import _is_namedtuple if _is_namedtuple(obj): - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] return [type(obj)(*args) for args in zip(*map(to_map, obj))] if isinstance(obj, tuple) and len(obj) > 0: - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] return list(zip(*map(to_map, obj))) if isinstance(obj, list) and len(obj) > 0: - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] return [list(i) for i in zip(*map(to_map, obj))] if isinstance(obj, dict) and len(obj) > 0: - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] return [type(obj)(i) for i in zip(*map(to_map, obj.items()))] return [obj] diff --git a/torch/distributions/kl.py b/torch/distributions/kl.py index 85932828d21af..43991b0ac265a 100644 --- a/torch/distributions/kl.py +++ b/torch/distributions/kl.py @@ -970,4 +970,5 @@ def _add_kl_info(): ) kl_info = "\n\t".join(rows) if kl_divergence.__doc__: + # pyrefly: ignore [missing-attribute] kl_divergence.__doc__ += kl_info diff --git a/torch/export/experimental/_utils.py b/torch/export/experimental/_utils.py index 9684e3530f89f..45c26f7de517d 100644 --- a/torch/export/experimental/_utils.py +++ b/torch/export/experimental/_utils.py @@ -196,7 +196,7 @@ def _get_make_file(package_name: str, model_names: list[str], device_type: str) "cmake_minimum_required(VERSION 3.10)", "project(TestProject)", "", - "set(CMAKE_CXX_STANDARD 17)", + "set(CMAKE_CXX_STANDARD 20)", "", ] ) diff --git a/torch/export/pt2_archive/_package.py b/torch/export/pt2_archive/_package.py index f8f413f9cfe60..2de31f40b20cb 100644 --- a/torch/export/pt2_archive/_package.py +++ b/torch/export/pt2_archive/_package.py @@ -603,6 +603,8 @@ def _package_exported_programs( ep, opset_version, pickle_protocol, + serialize_state_dict=False, + serialize_constants=False, ) archive_writer.write_bytes( diff --git a/torch/export/unflatten.py b/torch/export/unflatten.py index c76c52c37e152..d67c29ed46db4 100644 --- a/torch/export/unflatten.py +++ b/torch/export/unflatten.py @@ -1650,7 +1650,10 @@ def _reorder_submodules( fqn = prefix + name _reorder_submodules(child, fqn_order, prefix=fqn.split("@")[0] + ".") delattr(parent, name) - children.append((fqn_order[fqn], name, child)) + base_fqn = fqn.split("@")[0] + children.append( + (fqn_order.get(fqn, fqn_order.get(base_fqn, len(fqn_order))), name, child) + ) children.sort(key=operator.itemgetter(0)) for _, name, child in children: parent.register_module(name, child) diff --git a/torch/func/_random.py b/torch/func/_random.py new file mode 100644 index 0000000000000..a0b6fe6b5ab1f --- /dev/null +++ b/torch/func/_random.py @@ -0,0 +1,263 @@ +"""Stateless PRNG APIs. + +These are experimental and subject to change without notice. +Access via ``torch.func._random``. +""" + +from collections.abc import Sequence + +import torch + + +def key( + seed: int, impl: str = "philox4x32-10", device: torch.device | None = None +) -> torch.Tensor: + r"""Create a PRNG key from a seed. + + A key is a tensor that encodes the state needed to deterministically + produce random values. Keys are consumed by generation functions to produce + reproducible random tensors without any global state. The internal + representation of the key depends on the chosen PRNG algorithm. + + Args: + seed (int): The seed value for the PRNG. + impl (str): PRNG algorithm. Currently only ``"philox4x32-10"`` is + supported. + device (:class:`torch.device`, optional): The desired device for the + returned key. Default: ``cpu``. + + Returns: + A tensor representing the PRNG key. + + .. note:: + + For the ``"philox4x32-10"`` algorithm, the key is a uint64 tensor of + shape ``(2,)`` encoding a ``(seed, offset)`` pair. The offset determines + the starting position in the Philox output stream. + + Example:: + + >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP + """ + if impl != "philox4x32-10": + raise NotImplementedError(f"key() does not support PRNG impl '{impl}'") + + # (seed, offset) + return torch.tensor([seed, 0], dtype=torch.uint64, device=device) + + +def split(key: torch.Tensor, num: int = 2) -> torch.Tensor: + r"""Split a PRNG key into ``num`` new independent keys. + + Each returned key produces a different, deterministic random sequence. + This is the primary mechanism for deriving multiple independent keys from + a single parent key without mutating any state. + + Supports batched keys: if ``key`` has shape ``(*batch, K)``, each key in the + batch is split independently and the result has shape ``(num, *batch, K)``. + + Args: + key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or + :func:`fold_in`. + num (int): Number of keys to produce. Default: ``2``. + + Returns: + A tensor of shape ``(num, *key.shape)`` containing the derived keys. + + Example:: + + >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP + >>> k1, k2 = torch.func._random.split(key) # doctest: +SKIP + """ + return torch.ops.aten._philox_key_split(key, num) + + +def fold_in(key: torch.Tensor, data: int) -> torch.Tensor: + r"""Deterministically derive a new key by folding in an integer. + + Equivalent to ``split(key, data + 1)[data]``, but more efficient when + only a single derived key is needed. Useful for associating a key with + a loop iteration, layer index, or other integer identifier. + + Supports batched keys: if ``key`` has shape ``(*batch, K)``, each key in + the batch is folded independently. + + Args: + key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or + :func:`fold_in`. + data (int): An integer to fold into the key, interpreted as uint64. + + Returns: + A new key tensor with the same shape as ``key``. + + Example:: + + >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP + >>> k0 = torch.func._random.fold_in(key, 0) # doctest: +SKIP + >>> k1 = torch.func._random.fold_in(key, 1) # doctest: +SKIP + >>> # Equivalent to split: + >>> keys = torch.func._random.split(key, 2) # doctest: +SKIP + >>> assert torch.equal(k0, keys[0]) # doctest: +SKIP + >>> assert torch.equal(k1, keys[1]) # doctest: +SKIP + """ + return torch.ops.aten._philox_key_fold_in(key, data) + + +def normal_( + key: torch.Tensor, + result: torch.Tensor, + *, + mean: float = 0.0, + std: float = 1.0, +) -> torch.Tensor: + r"""Fill ``result`` in-place with normal random values from a PRNG key. + + The values are drawn from a normal distribution with the specified ``mean`` + and ``std``. The output is fully determined by the key, so calling with the + same key always produces the same result. + + Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading + dimensions of ``result`` must be broadcastable with ``*batch`` and each key + independently generates its slice of the output. + + Args: + key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or + :func:`fold_in`. + result (Tensor): The output tensor to fill in-place. + mean (float): Mean of the normal distribution. Default: ``0.0``. + std (float): Standard deviation of the normal distribution. Default: ``1.0``. + + Returns: + ``result``, filled with normal random values. + + Example:: + + >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP + >>> result = torch.empty(1000, device="cuda") # doctest: +SKIP + >>> torch.func._random.normal_(key, result) # doctest: +SKIP + """ + return torch.ops.aten._philox_normal_(result, key, mean, std) + + +def normal( + key: torch.Tensor, + *shape: tuple[int, ...], + mean: float = 0.0, + std: float = 1.0, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + r"""Generate normally distributed random values from a PRNG key. + + Produces a tensor of the given shape filled with values drawn from a normal + distribution with the specified ``mean`` and ``std``. The output is fully + determined by the key, so calling with the same key always returns the same + result. The output is placed on the same device as ``key``. + + Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading + dimensions of ``shape`` must be broadcastable with ``*batch`` and each key + independently generates its slice of the output. + + Args: + key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or + :func:`fold_in`. + *shape (int): The desired output shape. + mean (float): Mean of the normal distribution. Default: ``0.0``. + std (float): Standard deviation of the normal distribution. Default: ``1.0``. + dtype (:class:`torch.dtype`, optional): The desired dtype. Default: ``torch.float32``. + + Returns: + A tensor of the given shape filled with normal random values. + + Example:: + + >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP + >>> torch.func._random.normal(key, (1000,)) # doctest: +SKIP + """ + if len(shape) == 1 and isinstance(shape[0], Sequence): + # pyrefly: ignore [bad-argument-type] + shape = tuple(shape[0]) + if dtype is None: + dtype = torch.float32 + # pyrefly: ignore [no-matching-overload] + result = torch.empty(shape, dtype=dtype, device=key.device) + return normal_(key, result, mean=mean, std=std) + + +def uniform_( + key: torch.Tensor, + result: torch.Tensor, + *, + low: float = 0.0, + high: float = 1.0, +) -> torch.Tensor: + r"""Fill ``result`` in-place with uniform random values from a PRNG key. + + The values are drawn uniformly from the interval ``[low, high)``. The output + is fully determined by the key, so calling with the same key always produces + the same result. + + Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading + dimensions of ``result`` must be broadcastable with ``*batch`` and each key + independently generates its slice of the output. + + Args: + key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or + :func:`fold_in`. + result (Tensor): The output tensor to fill in-place. + low (float): Lower bound (inclusive) of the uniform distribution. Default: ``0.0``. + high (float): Upper bound (exclusive) of the uniform distribution. Default: ``1.0``. + + Returns: + ``result``, filled with uniform random values. + + Example:: + + >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP + >>> result = torch.empty(1000, device="cuda") # doctest: +SKIP + >>> torch.func._random.uniform_(key, result) # doctest: +SKIP + """ + return torch.ops.aten._philox_uniform_(result, key, low, high) + + +def uniform( + key: torch.Tensor, + *shape: tuple[int, ...], + low: float = 0.0, + high: float = 1.0, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + r"""Generate uniformly distributed random values from a PRNG key. + + Produces a tensor of the given shape filled with values drawn uniformly + from the interval ``[low, high)``. The output is fully determined by the + key, so calling with the same key always returns the same result. The output + is placed on the same device as ``key``. + + Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading + dimensions of ``shape`` must be broadcastable with ``*batch`` and each key + independently generates its slice of the output. + + Args: + key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or + :func:`fold_in`. + *shape (int): The desired output shape. + low (float): Lower bound (inclusive) of the uniform distribution. Default: ``0.0``. + high (float): Upper bound (exclusive) of the uniform distribution. Default: ``1.0``. + dtype (:class:`torch.dtype`, optional): The desired dtype. Default: ``torch.float32``. + + Returns: + A tensor of the given shape filled with uniform random values. + + Example:: + + >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP + >>> torch.func._random.uniform(key, (1000,)) # doctest: +SKIP + """ + if len(shape) == 1 and isinstance(shape[0], Sequence): + # pyrefly: ignore [bad-argument-type] + shape = tuple(shape[0]) + if dtype is None: + dtype = torch.float32 + # pyrefly: ignore [no-matching-overload] + result = torch.empty(shape, dtype=dtype, device=key.device) + return uniform_(key, result, low=low, high=high) diff --git a/torch/fx/_compatibility.py b/torch/fx/_compatibility.py index c07dd1b51bc05..265369cb84c21 100644 --- a/torch/fx/_compatibility.py +++ b/torch/fx/_compatibility.py @@ -16,6 +16,7 @@ def compatibility(is_backward_compatible: bool) -> Callable[[_T], _T]: def mark_back_compat(fn: _T) -> _T: docstring = textwrap.dedent(getattr(fn, "__doc__", None) or "") docstring += """ + .. note:: Backwards-compatibility for this API is guaranteed. """ @@ -30,6 +31,7 @@ def mark_back_compat(fn: _T) -> _T: def mark_not_back_compat(fn: _T) -> _T: docstring = textwrap.dedent(getattr(fn, "__doc__", None) or "") docstring += """ + .. warning:: This API is experimental and is *NOT* backward-compatible. """ diff --git a/torch/fx/_graph_pickler.py b/torch/fx/_graph_pickler.py index 47ac510c28871..da0d4f732b7f2 100644 --- a/torch/fx/_graph_pickler.py +++ b/torch/fx/_graph_pickler.py @@ -4,8 +4,9 @@ import io import itertools import pickle +import weakref from abc import abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Generator from typing import Any, NewType, TypeVar from typing_extensions import override, Self @@ -20,6 +21,7 @@ import torch.utils._pytree as pytree from torch._guards import TracingContext from torch._inductor.standalone_compile import AOTCompiledArtifact +from torch._library.fake_class_registry import FakeScriptObject from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode, Tensor from torch._subclasses.meta_utils import ( MetaConverter, @@ -73,8 +75,16 @@ def _unpickle_as_none() -> None: return None +def _unpickle_as_weakref(referent: object) -> weakref.ref[object]: + return weakref.ref(referent) + + +def _unpickle_as_dead_weakref() -> Callable[[], None]: + return lambda: None + + @contextlib.contextmanager -def patch_pytree_map_over_slice(): +def patch_pytree_map_over_slice() -> Generator[None]: if slice in pytree.SUPPORTED_NODES: yield return @@ -147,6 +157,20 @@ def reducer_override( return _SymNodePickleData.reduce_helper(self, obj) elif isinstance(obj, torch._guards.TracingContext): return _TracingContextPickleData.reduce_helper(self, obj) + elif isinstance(obj, FakeScriptObject): + # FakeScriptObjects wrap opaque traced objects (e.g. DeviceMesh, + # ProcessGroup) that can't be default-pickled. Reduce to None + # since they aren't meaningful after deserialization. + return (_unpickle_as_none, ()) + elif isinstance(obj, weakref.ref): + # Serialize weakrefs properly: if the referent is alive, + # serialize it and reconstruct the weakref on unpickle. + # If the referent is dead, unpickle as a dead-weakref-like callable. + referent = obj() + if referent is not None: + return (_unpickle_as_weakref, (referent,)) + else: + return (_unpickle_as_dead_weakref, ()) else: # We should never get a raw Node! if isinstance(obj, torch.fx.Node): diff --git a/torch/fx/_lazy_graph_module.py b/torch/fx/_lazy_graph_module.py index 7808dec3cee13..20a7837a66618 100644 --- a/torch/fx/_lazy_graph_module.py +++ b/torch/fx/_lazy_graph_module.py @@ -1,5 +1,6 @@ -# mypy: allow-untyped-defs +from collections.abc import Iterator from contextlib import contextmanager +from typing import Any, TYPE_CHECKING from torch.fx.graph_module import ( _format_import_block, @@ -12,13 +13,17 @@ from ._compatibility import compatibility +if TYPE_CHECKING: + from torch.fx.graph import PythonCode + + _use_lazy_graph_module_flag = False _force_skip_lazy_graph_module_flag = False @compatibility(is_backward_compatible=False) @contextmanager -def _force_skip_lazy_graph_module(): +def _force_skip_lazy_graph_module() -> Iterator[None]: """ Skip using lazy graph module disregarding the setting of _use_lazy_graph_module. Use to skip _LazyGraphModule when testing inductor torchscript related backend. @@ -37,7 +42,7 @@ def _force_skip_lazy_graph_module(): @compatibility(is_backward_compatible=False) @contextmanager -def _use_lazy_graph_module(should_use: bool): +def _use_lazy_graph_module(should_use: bool) -> Iterator[None]: try: global _use_lazy_graph_module_flag prior = _use_lazy_graph_module_flag @@ -50,11 +55,13 @@ def _use_lazy_graph_module(should_use: bool): @compatibility(is_backward_compatible=False) -def _get_graph_module_cls(): +def _get_graph_module_cls() -> type[GraphModule]: return _LazyGraphModule if _use_lazy_graph_module_flag else GraphModule -def _make_graph_module(*args, graph_module_cls=None, **kwargs): +def _make_graph_module( + *args: Any, graph_module_cls: type[GraphModule] | None = None, **kwargs: Any +) -> GraphModule: if graph_module_cls is None: graph_module_cls = _get_graph_module_cls() @@ -88,14 +95,14 @@ class _LazyGraphModule(GraphModule): """ @classmethod - def from_graphmodule(cls, gm: GraphModule): + def from_graphmodule(cls, gm: GraphModule) -> GraphModule: if isinstance(gm, _LazyGraphModule): return gm else: return _LazyGraphModule(gm, gm.graph) @staticmethod - def force_recompile(gm): + def force_recompile(gm: GraphModule) -> None: """ Sometimes we need force a recompile as a workaround - we want to do the real recompilation before symbolic_trace to avoid error: @@ -104,15 +111,15 @@ def force_recompile(gm): if isinstance(gm, _LazyGraphModule): gm.real_recompile() - def real_recompile(self): + def real_recompile(self) -> None: if self._needs_recompile(): self._real_recompile() @classmethod - def _needs_recompile(cls): + def _needs_recompile(cls) -> bool: return cls.forward is cls._lazy_forward - def _lazy_forward(self, *args, **kwargs): + def _lazy_forward(self, *args: Any, **kwargs: Any) -> Any: # Call self.real_recompile() rather than self._real_recompile() here. # The _lazy_forward method may be saved and call repeatedly. # Calling self.real_recompile can make sure we skip recompilation if @@ -128,7 +135,9 @@ def _lazy_forward(self, *args, **kwargs): forward = _lazy_forward - def __reduce_package__(self, exporter: PackageExporter): + def __reduce_package__( + self, exporter: PackageExporter + ) -> tuple[Any, tuple[Any, str]]: """ Follow GraphModule.__reduce__ but call 'self._real_recompile' rather than 'self.recompile' since for a _LazyGraphModule, self.recompile just @@ -148,7 +157,7 @@ def __reduce_package__(self, exporter: PackageExporter): (dict_without_graph, generated_module_name), ) - def __reduce__(self): + def __reduce__(self) -> tuple[Any, tuple[Any, str]]: """ Follow GraphModule.__reduce__ but call 'self._real_recompile' rather than 'self.recompile' since for a _LazyGraphModule, self.recompile just @@ -160,11 +169,11 @@ def __reduce__(self): del dict_without_graph["_graph"] return (reduce_graph_module, (dict_without_graph, import_block)) - def _real_recompile(self): + def _real_recompile(self) -> "PythonCode": return super().recompile() @classmethod - def recompile(cls): + def recompile(cls) -> None: # pyrefly: ignore[bad-override] cls.forward = cls._lazy_forward @property diff --git a/torch/fx/_pytree.py b/torch/fx/_pytree.py index 5b18c815ca312..4d923d528e710 100644 --- a/torch/fx/_pytree.py +++ b/torch/fx/_pytree.py @@ -7,7 +7,7 @@ from torch.utils._pytree import PyTree, tree_flatten, TreeSpec -FlattenFnSpec = Callable[[PyTree, TreeSpec], list] +FlattenFnSpec = Callable[[PyTree, TreeSpec], list[Any]] FlattenFnExactMatchSpec = Callable[[PyTree, TreeSpec], bool] # Keep deprecated alias for backward compatibility @@ -52,7 +52,7 @@ def tree_flatten_spec( if spec.type in SUPPORTED_NODES: flatten_fn_spec = SUPPORTED_NODES[spec.type] child_pytrees = flatten_fn_spec(pytree, spec) - result = [] + result: list[Any] = [] for child, child_spec in zip(child_pytrees, spec.children()): flat = tree_flatten_spec(child, child_spec) result += flat diff --git a/torch/fx/_symbolic_trace.py b/torch/fx/_symbolic_trace.py index 893225b3836ef..3a5541811cb6e 100644 --- a/torch/fx/_symbolic_trace.py +++ b/torch/fx/_symbolic_trace.py @@ -1,4 +1,3 @@ -# mypy: allow-untyped-defs import builtins import collections import contextlib @@ -9,10 +8,10 @@ import math import os import warnings -from collections.abc import Callable +from collections.abc import Callable, Iterable, Iterator from itertools import chain -from types import CodeType, FunctionType, ModuleType -from typing import Any, get_args, NamedTuple, TypeAlias +from types import CodeType, FunctionType, ModuleType, TracebackType +from typing import Any, get_args, NamedTuple, overload, ParamSpec, TypeAlias, TypeVar import torch import torch.utils._pytree as pytree @@ -32,9 +31,13 @@ HAS_VARSTUFF = inspect.CO_VARARGS | inspect.CO_VARKEYWORDS +_F = TypeVar("_F", bound=Callable[..., Any]) +_P = ParamSpec("_P") +_T = TypeVar("_T") + # These need to run in global scope to handle nested calls correctly -_orig_module_call: Callable = torch.nn.Module.__call__ -_orig_module_getattr: Callable = torch.nn.Module.__getattr__ +_orig_module_call: Callable[..., Any] = torch.nn.Module.__call__ +_orig_module_getattr: Callable[..., Any] = torch.nn.Module.__getattr__ _proxyable_classes: dict[type, None] = {} @@ -49,7 +52,7 @@ # We only want to print this once to avoid flooding logs @functools.lru_cache -def is_fx_tracing_warning(): +def is_fx_tracing_warning() -> None: log.warning( "is_fx_tracing will return true for both fx.symbolic_trace and " "torch.export. Please use " @@ -58,12 +61,12 @@ def is_fx_tracing_warning(): ) -def is_fx_tracing(): +def is_fx_tracing() -> bool: is_fx_tracing_warning() return _is_fx_tracing_flag -def is_fx_symbolic_tracing(): +def is_fx_symbolic_tracing() -> bool: return _is_fx_tracing_flag and not torch.compiler.is_compiling() @@ -116,20 +119,22 @@ def forward(self, x : __main___TensorPair, y : torch.Tensor): tracing. """ - def __init__(cls, name, bases, attrs): + def __init__( + cls, name: str, bases: tuple[type, ...], attrs: dict[str, Any] + ) -> None: _proxyable_classes.setdefault(cls) super().__init__(name, bases, attrs) - def __call__(cls, *args, **kwargs): + def __call__(cls, *args: Any, **kwargs: Any) -> Any: instance = cls.__new__(cls) # type: ignore[call-overload] if not is_fx_tracing(): cls.__init__(instance, *args, **kwargs) # type: ignore[misc] return instance - found_proxies = [] + found_proxies: list[Proxy] = [] - def check_proxy(a): + def check_proxy(a: object) -> None: if isinstance(a, Proxy): found_proxies.append(a) @@ -147,7 +152,7 @@ def check_proxy(a): def _patch_function(fn: FunctionType, nargs: int) -> FunctionType: co = fn.__code__ co_flags = co.co_flags & ~HAS_VARSTUFF - co_args: tuple + co_args: tuple[Any, ...] if hasattr(co, "co_qualname"): # Python-3.11+ code signature co_args = ( @@ -223,7 +228,7 @@ class PHBase: Object representing an input placeholder to `concrete_args` """ - def __repr__(self): + def __repr__(self) -> str: return "PH" @@ -236,14 +241,14 @@ class PHWithMeta(PHBase): Object representing an input placeholder to `concrete_args` """ - def __init__(self, ph_key: str | None = None): + def __init__(self, ph_key: str | None = None) -> None: super().__init__() # Provide a hey for user to identify placeholder node during analysis self.ph_key = ph_key -def _transfer_attrs(fr, to): +def _transfer_attrs(fr: object, to: object) -> None: for attr_name in dir(fr): attr_val = getattr(fr, attr_name) if ( @@ -279,7 +284,7 @@ class Tracer(TracerBase): def __init__( self, autowrap_modules: tuple[ModuleType] = (math,), - autowrap_functions: tuple[Callable, ...] = (), + autowrap_functions: tuple[Callable[..., Any], ...] = (), param_shapes_constant: bool = False, ) -> None: # This method's signature is overridden by the first line of this class' @@ -581,7 +586,9 @@ def call_module( return ret_val @compatibility(is_backward_compatible=False) - def getattr(self, attr: str, attr_val: Any, parameter_proxy_cache: dict[str, Any]): + def getattr( + self, attr: str, attr_val: Any, parameter_proxy_cache: dict[str, Proxy] + ) -> Any: """ Method that specifies the behavior of this ``Tracer`` when we call getattr on a call to an ``nn.Module`` instance. @@ -605,8 +612,10 @@ def getattr(self, attr: str, attr_val: Any, parameter_proxy_cache: dict[str, Any """ def maybe_get_proxy_for_attr( - attr_val, collection_to_search, parameter_proxy_cache - ): + attr_val: Any, + collection_to_search: Iterable[tuple[str, Any]], + parameter_proxy_cache: dict[str, Proxy], + ) -> Proxy | None: for n, p in collection_to_search: if attr_val is p: if n not in parameter_proxy_cache: @@ -646,7 +655,12 @@ def maybe_get_proxy_for_attr( # This method will be refactored @compatibility(is_backward_compatible=False) - def create_args_for_root(self, root_fn, is_module, concrete_args=None): + def create_args_for_root( + self, + root_fn: Callable[..., Any], + is_module: bool, + concrete_args: dict[str, Any] | tuple[Any, ...] | None = None, + ) -> tuple[Any, list[Any]]: """ Create ``placeholder`` nodes corresponding to the signature of the ``root`` Module. This method introspects root's signature and emits those @@ -708,7 +722,7 @@ def create_args_for_root(self, root_fn, is_module, concrete_args=None): ) concrete_args = dict(zip(arg_names, concrete_args)) - def proxy_placeholder(name): + def proxy_placeholder(name: str) -> Any: return self._proxy_placeholder(name, concrete_args, sig, fn_for_analysis) args.extend(proxy_placeholder(names) for names in arg_names) @@ -730,7 +744,9 @@ def proxy_placeholder(name): _PyTreeInfo(orig_args[:total_args], in_spec, None) ) - def flatten_fn(*args): + # TODO: annotate return type. inspect.get_annotations(flatten_fn) + # leaks the return annotation into the generated forward() code. + def flatten_fn(*args: Any): # pyrefly: ignore[unannotated-parameter] tree_args = pytree.tree_unflatten(list(args), in_spec) tree_out = root_fn(*tree_args) out_args, out_spec = pytree.tree_flatten(tree_out) @@ -787,7 +803,9 @@ def trace( # without this. from torch.fx._lazy_graph_module import _LazyGraphModule - _LazyGraphModule.force_recompile(root) + _LazyGraphModule.force_recompile( + root # pyrefly: ignore[bad-argument-type] + ) self.root = root @@ -823,7 +841,9 @@ def trace( str, ] = {} - def collect_tensor_attrs(m: torch.nn.Module, prefix_atoms: list[str]): + def collect_tensor_attrs( + m: torch.nn.Module, prefix_atoms: list[str] + ) -> None: for k, v in m.__dict__.items(): if isinstance(v, _constant_attribute_types): self.tensor_attrs[v] = ".".join(prefix_atoms + [k]) @@ -847,13 +867,15 @@ def collect_tensor_attrs(m: torch.nn.Module, prefix_atoms: list[str]): # Method dispatch on parameters is not recorded unless it's directly used. # Thus, we need to insert a proxy when __getattr__ requests a parameter. @functools.wraps(_orig_module_getattr) - def module_getattr_wrapper(mod, attr): + def module_getattr_wrapper(mod: torch.nn.Module, attr: str) -> Any: attr_val = _orig_module_getattr(mod, attr) return self.getattr(attr, attr_val, parameter_proxy_cache) @functools.wraps(_orig_module_call) - def module_call_wrapper(mod, *args, **kwargs): - def forward(*args, **kwargs): + def module_call_wrapper( + mod: torch.nn.Module, *args: Any, **kwargs: Any + ) -> Any: + def forward(*args: Any, **kwargs: Any) -> Any: return _orig_module_call(mod, *args, **kwargs) _autowrap_check( @@ -907,7 +929,7 @@ def forward(*args, **kwargs): _is_fx_tracing_flag = old_is_fx_tracing_flag return self.graph - def __deepcopy__(self, memo): + def __deepcopy__(self, memo: dict[int, Any]) -> "Tracer": # _autowrap_search contains modules, which cannot be deepcopied. new_tracer = Tracer.__new__(Tracer) @@ -921,11 +943,17 @@ def __deepcopy__(self, memo): return new_tracer - def _proxy_placeholder(self, name, concrete_args, sig, fn_for_analysis): + def _proxy_placeholder( + self, + name: str, + concrete_args: dict[str, Any] | None, + sig: inspect.Signature, + fn_for_analysis: Callable[..., Any], + ) -> Any: if concrete_args is not None and name in concrete_args: cnt = 0 - def replace_ph(x): + def replace_ph(x: object) -> object: nonlocal cnt cnt += 1 param = sig.parameters[name] @@ -993,7 +1021,7 @@ def replace_ph(x): # the purposes of the wrap() API. # We key by the globals dict id and function name to ensure we're wrapping a given # function only once. -_wrapped_fns_to_patch: dict[tuple[int, str], dict] = {} +_wrapped_fns_to_patch: dict[tuple[int, str], dict[str, Any]] = {} # List of methods on classes to wrap (class type, function name) # this currently only works for Tensor.* methods that aren't traced properly @@ -1008,25 +1036,26 @@ def replace_ph(x): _wrapped_methods_to_patch.append((torch.Tensor, "__getitem__")) -def _find_proxy(*objects_to_search): +def _find_proxy(*objects_to_search: object) -> Proxy | None: """ Recursively search a data structure for a Proxy() and return it, return None if not found. """ proxy = None - def find_proxy(x): + def find_proxy(x: object) -> None: nonlocal proxy if isinstance(x, Proxy): proxy = x + # pyrefly: ignore[bad-specialization] map_aggregate(objects_to_search, find_proxy) return proxy -def _create_wrapped_func(orig_fn): +def _create_wrapped_func(orig_fn: Callable[_P, _T]) -> Callable[_P, _T]: @functools.wraps(orig_fn) - def wrapped(*args, **kwargs): + def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> Any: """ Given an closed-over ``orig_function`` to invoke, search the args and kwargs for a Proxy object. If there is one, emit a ``call_function`` node to preserve the @@ -1045,11 +1074,11 @@ def wrapped(*args, **kwargs): return wrapped -def _create_wrapped_method(cls, name): +def _create_wrapped_method(cls: type, name: str) -> Callable[..., Any]: orig_fn = getattr(cls, name) @functools.wraps(orig_fn) - def wrapped(*args, **kwargs): + def wrapped(*args: Any, **kwargs: Any) -> Any: """ Search the args and kwargs for a Proxy object. If there is one, emit a ``call_method`` node to preserve the call to this method @@ -1065,39 +1094,39 @@ def wrapped(*args, **kwargs): class _PatchedFn(NamedTuple): - frame_dict: Any + frame_dict: Any # dict[str, Any] for SetItem/Del, type for SetAttr fn_name: str - orig_fn: Any - new_fn: Any + orig_fn: Callable[..., Any] | None + new_fn: Callable[..., Any] - def revert(self): + def revert(self) -> None: raise NotImplementedError - def patch(self): + def patch(self) -> None: raise NotImplementedError class _PatchedFnSetItem(_PatchedFn): - def revert(self): + def revert(self) -> None: self.frame_dict[self.fn_name] = self.orig_fn - def patch(self): + def patch(self) -> None: self.frame_dict[self.fn_name] = self.new_fn class _PatchedFnDel(_PatchedFn): - def revert(self): + def revert(self) -> None: del self.frame_dict[self.fn_name] - def patch(self): + def patch(self) -> None: self.frame_dict[self.fn_name] = self.new_fn class _PatchedFnSetAttr(_PatchedFn): - def revert(self): + def revert(self) -> None: setattr(self.frame_dict, self.fn_name, self.orig_fn) - def patch(self): + def patch(self) -> None: setattr(self.frame_dict, self.fn_name, self.new_fn) @@ -1111,9 +1140,9 @@ def patch( self, frame_dict: dict[str, Any], name: str, - new_fn: Callable, + new_fn: Callable[..., Any], deduplicate: bool = True, - ): + ) -> None: """ Replace frame_dict[name] with new_fn until we exit the context manager. """ @@ -1130,8 +1159,8 @@ def patch( self.patches_made[-1].patch() def patch_method( - self, cls: type, name: str, new_fn: Callable, deduplicate: bool = True - ): + self, cls: type, name: str, new_fn: Callable[..., Any], deduplicate: bool = True + ) -> None: """ Replace object_or_dict.name with new_fn until we exit the context manager. """ @@ -1142,7 +1171,7 @@ def patch_method( self.patches_made.append(_PatchedFnSetAttr(cls, name, orig_fn, new_fn)) self.patches_made[-1].patch() - def visit_once(self, thing: Any): + def visit_once(self, thing: object) -> bool: """Return True on the first call to with thing, otherwise false""" idx = id(thing) if idx in self.visited: @@ -1150,7 +1179,7 @@ def visit_once(self, thing: Any): self.visited.add(idx) return True - def revert_all_patches(self): + def revert_all_patches(self) -> list[_PatchedFn]: """ Remove all the stored patcheds. It doesn't modify patches_made. """ @@ -1158,7 +1187,7 @@ def revert_all_patches(self): patch.revert() return self.patches_made - def reapply_all_patches(self): + def reapply_all_patches(self) -> list[_PatchedFn]: """ Patch all the stored patcheds. It doesn't modify patches_made. """ @@ -1166,10 +1195,15 @@ def reapply_all_patches(self): patch.patch() return self.patches_made - def __enter__(self): + def __enter__(self) -> "_Patcher": return self - def __exit__(self, exc_type, exc_val, exc_tb): + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: """ Undo all the changes made via self.patch() and self.patch_method() """ @@ -1183,7 +1217,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): @contextlib.contextmanager -def _new_patcher(): +def _new_patcher() -> Iterator[_Patcher]: global CURRENT_PATCHER prior_patcher = CURRENT_PATCHER try: @@ -1198,7 +1232,7 @@ def _new_patcher(): @contextlib.contextmanager -def _maybe_revert_all_patches(): +def _maybe_revert_all_patches() -> Iterator[None]: current_patcher = CURRENT_PATCHER patches_made = None patches_removed = None @@ -1215,7 +1249,7 @@ def _maybe_revert_all_patches(): ) -def _patch_wrapped_functions(patcher: _Patcher): +def _patch_wrapped_functions(patcher: _Patcher) -> None: """ Go through ``_wrapped_fn_patch_table`` and, for each frame object, wrap the listed global functions in the `_create_wrapped_func` wrapper. @@ -1233,7 +1267,7 @@ def _patch_wrapped_functions(patcher: _Patcher): def _autowrap_check( patcher: _Patcher, frame_dict: dict[str, Any], function_ids: set[int] -): +) -> None: """ Some methods, like `math.sqrt` are common enough we want to automatically wrap them as we see them. This method searches a scope for them and patches them if found. @@ -1248,8 +1282,12 @@ def _autowrap_check( patcher.patch(frame_dict, name, _create_wrapped_func(value)) +@overload +def wrap(fn_or_name: _F) -> _F: ... +@overload +def wrap(fn_or_name: str) -> str: ... @compatibility(is_backward_compatible=True) -def wrap(fn_or_name: str | Callable): +def wrap(fn_or_name: str | Callable[..., Any]) -> str | Callable[..., Any]: """ This function can be called at module-level scope to register fn_or_name as a "leaf function". A "leaf function" will be preserved as a CallFunction node in the FX trace instead of being @@ -1382,6 +1420,6 @@ def f(x): @wrap -def _assert_is_none(value, msg): +def _assert_is_none(value: object, msg: str) -> None: if value is not None: raise AssertionError(msg) diff --git a/torch/fx/_utils.py b/torch/fx/_utils.py index c27d7d113ffd8..c58c36502b68d 100644 --- a/torch/fx/_utils.py +++ b/torch/fx/_utils.py @@ -1,16 +1,18 @@ -# mypy: allow-untyped-defs import sys +from typing import Any import torch from torch._logging import LazyString -def lazy_format_graph_code(name, gm, maybe_id=None, **kwargs): +def lazy_format_graph_code( + name: str, gm: torch.fx.GraphModule, maybe_id: int | None = None, **kwargs: Any +) -> LazyString: """ Returns a LazyString that formats the graph code. """ - def format_name(): + def format_name() -> str: if maybe_id is not None: return f"{name} {maybe_id}" else: @@ -35,14 +37,14 @@ def format_name(): ) -def _format_graph_code(name, filename, graph_str): +def _format_graph_code(name: str, filename: str, graph_str: str) -> str: """ Returns a string that formats the graph code. """ return f"TRACED GRAPH\n {name} {filename} {graph_str}\n" -def first_call_function_nn_module_stack(graph: torch.fx.Graph) -> dict | None: +def first_call_function_nn_module_stack(graph: torch.fx.Graph) -> dict[str, Any] | None: """ Returns the nn_module_stack of the first call_function node. """ @@ -52,14 +54,15 @@ def first_call_function_nn_module_stack(graph: torch.fx.Graph) -> dict | None: return None -def get_node_context(node, num_nodes=2) -> str: +def get_node_context(node: torch.fx.Node, num_nodes: int = 2) -> str: """ Returns a string of the last num_nodes nodes in the graph. """ node_contexts = [] cur = node for _ in range(num_nodes): - node_contexts.append(cur.format_node()) + # cast to str to handle None return value + node_contexts.append(str(cur.format_node())) if cur.op == "root": break cur = cur.prev diff --git a/torch/fx/annotate.py b/torch/fx/annotate.py index b3c5056066251..fd90605d8e3b5 100644 --- a/torch/fx/annotate.py +++ b/torch/fx/annotate.py @@ -1,11 +1,15 @@ -# mypy: allow-untyped-defs +from typing import Any + from torch.fx.proxy import Proxy from ._compatibility import compatibility +__all__ = ["annotate"] + + @compatibility(is_backward_compatible=False) -def annotate(val, type): +def annotate(val: Any, type: type) -> Any: """ Annotates a Proxy object with a given type. diff --git a/torch/fx/experimental/_size_hinting.py b/torch/fx/experimental/_size_hinting.py index 80791e4dccf64..578054fb139cf 100644 --- a/torch/fx/experimental/_size_hinting.py +++ b/torch/fx/experimental/_size_hinting.py @@ -23,6 +23,11 @@ log = logging.getLogger(__name__) +# Maximum number of free symbols in an expression before we skip +# sympy.factor() in optimization_hint process for unbacked. +# Factoring polynomials with many variables is expensive. +SYMPY_FACTOR_MAX_FREE_SYMBOLS = 50 + if TYPE_CHECKING: from torch.fx.experimental.symbolic_shapes import ShapeEnv @@ -297,9 +302,7 @@ def _sub_unbacked_exprs(shape_env: ShapeEnv, expr: sympy.Expr) -> sympy.Expr: new_expr = expr.subs(replacements) if new_expr == expr: break - # Limit sympy.factor() to expressions with <= 200 free symbols, - # as factoring polynomials with many variables is expensive. - if len(new_expr.free_symbols) <= 200: + if len(new_expr.free_symbols) <= SYMPY_FACTOR_MAX_FREE_SYMBOLS: expr = sympy.factor(new_expr) else: expr = new_expr @@ -347,10 +350,14 @@ def _optimization_hint_base( fallback = unbacked_symint_fallback + # to have expanded (Identity free) expr stored in original + if isinstance(expr, sympy.Expr): + expr = expr.expand(identity=True) + original = expr # sympy.expand() doesn't work with boolean expressions like Or/And if isinstance(expr, sympy.Expr): - expr = sympy.expand(expr).xreplace(shape_env.replacements) + expr = expr.xreplace(shape_env.replacements) else: expr = sympy.sympify(expr).xreplace(shape_env.replacements) @@ -391,9 +398,10 @@ def _optimization_hint_base( if has_free_unbacked_symbols(expr): # Make sure to substitute with the factored version # e.g. 10*(s0 + u0) instead of 10*s0 + 10*u0 - # Limit sympy.factor() to expressions with <= 200 free symbols, - # as factoring polynomials with many variables is expensive. - if isinstance(original, sympy.Expr) and len(original.free_symbols) <= 200: + if ( + isinstance(original, sympy.Expr) + and len(original.free_symbols) <= SYMPY_FACTOR_MAX_FREE_SYMBOLS + ): original = sympy.factor(original) expr = _sub_unbacked_exprs(shape_env, original) diff --git a/torch/fx/experimental/const_fold.py b/torch/fx/experimental/const_fold.py index f44293e2242a7..d08e036fd797a 100644 --- a/torch/fx/experimental/const_fold.py +++ b/torch/fx/experimental/const_fold.py @@ -288,14 +288,27 @@ def mod_partition(node: torch.fx.Node): const_mod_name, non_const_mod_name = "submod_0", "submod_1" # Safely get submod_1 in case there are no non-const nodes - const_gm, non_const_gm = split.submod_0, getattr(split, non_const_mod_name, None) + const_gm = getattr(split, const_mod_name) + if not isinstance(const_gm, torch.fx.GraphModule): + raise AssertionError( + f"Expected GraphModule for {const_mod_name}, got {type(const_gm)}" + ) + non_const_mod = getattr(split, non_const_mod_name, None) + non_const_gm: torch.fx.GraphModule | None = None + if non_const_mod is not None: + if not isinstance(non_const_mod, torch.fx.GraphModule): + raise AssertionError( + f"Expected GraphModule for {non_const_mod_name}, got {type(non_const_mod)}" + ) + non_const_gm = non_const_mod # The module that a call_module node refers to gets copied to submodules during split. # The path to the module also gets inlined, i.e. mod.a.b -> mod_a_b. Here we need to # attach inlined modules to `split` as it's the owning module now. - for node in non_const_gm.graph.nodes if non_const_gm else []: - if node.op == "call_module": - setattr(split, node.target, getattr(non_const_gm, node.target)) + if non_const_gm is not None: + for node in non_const_gm.graph.nodes: + if node.op == "call_module": + setattr(split, node.target, getattr(non_const_gm, node.target)) for node in const_gm.graph.nodes: if node.op == "call_module": setattr(split, node.target, getattr(const_gm, node.target)) diff --git a/torch/fx/experimental/proxy_tensor.py b/torch/fx/experimental/proxy_tensor.py index a2f183b22a7b1..28787fd3057ba 100644 --- a/torch/fx/experimental/proxy_tensor.py +++ b/torch/fx/experimental/proxy_tensor.py @@ -1,4 +1,3 @@ -# mypy: allow-untyped-decorators # Copyright (c) Facebook, Inc. and its affiliates. # All rights reserved. # @@ -7,6 +6,7 @@ from __future__ import annotations +import contextvars import functools import inspect import logging @@ -14,7 +14,6 @@ import threading import typing import typing_extensions -import weakref from collections import defaultdict, OrderedDict from collections.abc import Callable, Generator, Mapping, Sequence from contextlib import _GeneratorContextManager, contextmanager, ExitStack, nullcontext @@ -22,6 +21,7 @@ from typing import ( Any, Concatenate, + Literal, overload, Protocol, TYPE_CHECKING, @@ -41,13 +41,15 @@ from torch._dispatch.python import enable_python_dispatcher from torch._library.fake_class_registry import FakeScriptObject from torch._library.opaque_object import ( + get_reconstruct_fn, + is_opaque_reference_type, is_opaque_value, is_opaque_value_type, - OpaqueType, + OpaqueBase, + should_hoist, ) from torch._logging import trace_structured -from torch._opaque_base import OpaqueBase -from torch._ops import HigherOrderOperator +from torch._ops import HigherOrderOperator, OpOverload from torch._subclasses.fake_impls import fast_detach from torch._subclasses.fake_tensor import ( FakeTensor, @@ -90,8 +92,6 @@ import sympy from torch._higher_order_ops.utils import FunctionalizeCtxWrapper - from torch._ops import OpOverload - from torch.fx._symbolic_trace import PHBase from torch.types import BoolLikeType, FloatLikeType, IntLikeType __all__ = [ @@ -109,6 +109,7 @@ ] _ProxyTracer = Union["PythonKeyTracer", "_GraphAppendingTracerEx"] +_TracingMode: TypeAlias = Literal["real", "fake", "symbolic"] _AnyScriptObject = (torch.ScriptObject, FakeScriptObject) _AnyScriptObjectType = torch.ScriptObject | FakeScriptObject @@ -119,7 +120,9 @@ log = logging.getLogger(__name__) not_implemented_log = torch._logging.getArtifactLogger(__name__, "not_implemented") -CURRENT_DECOMPOSITION_TABLE: Mapping[OpOverload, Callable] = {} +CURRENT_DECOMPOSITION_TABLE: contextvars.ContextVar[ + Mapping[OpOverload, Callable[..., Any]] +] = contextvars.ContextVar("CURRENT_DECOMPOSITION_TABLE") CONSTANT_NUMEL_LIMIT = 1 @@ -129,7 +132,6 @@ R = TypeVar("R") _Ts = TypeVarTuple("_Ts") -null_ctx_type = type(nullcontext) # We currently convert all SymInt to proxies before we use them. # This could plausibly be handled at the Dynamo level. pytree.register_pytree_node( @@ -160,15 +162,14 @@ def fake_signature(fn: Callable[_P, R], nargs: int) -> Callable[_P, R]: @contextmanager def decompose( - decomposition_table: Mapping[OpOverload, Callable] | None, -) -> Generator[Mapping[OpOverload, Callable], None, None]: - global CURRENT_DECOMPOSITION_TABLE - old_decomposition_table = CURRENT_DECOMPOSITION_TABLE - CURRENT_DECOMPOSITION_TABLE = decomposition_table or {} + decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None, +) -> Generator[Mapping[OpOverload, Callable[..., Any]], None, None]: + table = decomposition_table or {} + token = CURRENT_DECOMPOSITION_TABLE.set(table) try: - yield CURRENT_DECOMPOSITION_TABLE + yield table finally: - CURRENT_DECOMPOSITION_TABLE = old_decomposition_table + CURRENT_DECOMPOSITION_TABLE.reset(token) # ensure we cannot collide with other properties @@ -262,7 +263,7 @@ def set_proxy_slot(obj: Tensor, tracer: _ProxyTracer, proxy: _ProxyTensor) -> No @overload def set_proxy_slot( - obj: _AnyScriptObjectType | OpaqueType, tracer: _ProxyTracer, proxy: Proxy + obj: _AnyScriptObjectType | OpaqueBase, tracer: _ProxyTracer, proxy: Proxy ) -> None: ... @@ -273,7 +274,7 @@ def set_proxy_slot( def set_proxy_slot( - obj: PySymType | _AnyScriptObjectType | Tensor | OpaqueType, + obj: PySymType | _AnyScriptObjectType | Tensor | OpaqueBase, tracer: _ProxyTracer, proxy: object, ) -> None: @@ -290,7 +291,6 @@ def set_proxy_slot( ): tracer.tensor_tracker[obj] = proxy elif isinstance(obj, (_AnyScriptObject)) or is_opaque_value(obj): - # We DO want to clobber proxies, with a similar rationale as for tensors. if not isinstance(proxy, Proxy): raise AssertionError(f"Expected Proxy, got {type(proxy)}") # ScriptObject (actual C++ torchbind) uses _WeakHashRef-keyed tracker @@ -300,7 +300,30 @@ def set_proxy_slot( if isinstance(obj, torch.ScriptObject): tracer.script_object_tracker[obj] = proxy else: - tracer.opaque_tracker[obj] = proxy + # NB: Never clobber a pre-existing proxy for the same + # underlying real object. Multiple FakeScriptObject wrappers + # can share the same real_obj (e.g. primal vs tangent + # placeholders during joint graph tracing). We always keep the + # first proxy registered, with the same rationale as the + # symnode_tracker first-one-wins policy below: primals are + # registered first, so this avoids spurious tangent dependencies + # in forward outputs (which would break the partitioner). + real_obj = None + if isinstance(obj, FakeScriptObject): + try: + real_obj = object.__getattribute__(obj, "real_obj") + except AttributeError: + pass + + if real_obj is not None: + existing = tracer._opaque_real_obj_proxy.get(id(real_obj)) + if existing is not None: + tracer.opaque_tracker[obj] = existing + else: + tracer.opaque_tracker[obj] = proxy + tracer._opaque_real_obj_proxy[id(real_obj)] = proxy + else: + tracer.opaque_tracker[obj] = proxy else: # NB: Never clobber pre-existing proxy. Although the proxies # are in principle equivalent, when we do graph partitioning @@ -342,7 +365,7 @@ def has_proxy_slot(obj: Tensor, tracer: _ProxyTracer) -> bool: _PySymProxyType: TypeAlias = Thunk[Proxy] -_OpaqueObjectProxyType: TypeAlias = Thunk[Proxy] +_OpaqueObjectProxyType: TypeAlias = Proxy @overload @@ -410,7 +433,7 @@ def get_proxy_slot( @overload def get_proxy_slot( - obj: OpaqueType, + obj: OpaqueBase, tracer: _ProxyTracer, default: T, ) -> T | _OpaqueObjectProxyType: ... @@ -429,10 +452,10 @@ def get_proxy_slot( # the transform argument is handy if you need to extract a subfield from # the successfully looked up result (but NOT the default.) def get_proxy_slot( - obj: Tensor | _AnyScriptObjectType | PySymType | OpaqueType, + obj: Tensor | _AnyScriptObjectType | PySymType | OpaqueBase, tracer: _ProxyTracer, default: object = no_default, - transform: Callable = lambda x: x, + transform: Callable[..., Any] = lambda x: x, ) -> object: tracker: Any if isinstance(obj, Tensor): @@ -461,9 +484,16 @@ def get_proxy_slot( else: # Attempt to build it from first principles. _build_proxy_for_sym_expr(tracer, obj.node.expr, obj) - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] value = tracker.get(obj) + if value is None and isinstance(obj, FakeScriptObject): + # A new FakeScriptObject wrapping the same real_obj may have been + # created (e.g. output flattening in unwrap_tensor_subclasses calls + # maybe_to_fake_obj which always mints a fresh wrapper). Fall back + # to the real-object dedup map that set_proxy_slot maintains. + value = tracer._opaque_real_obj_proxy.get(id(obj.real_obj)) + if value is None: # We don't know this value - return the default. if isinstance(default, _NoDefault): @@ -638,7 +668,7 @@ def snapshot_fake(val: Tensor, include_real: bool = False) -> Tensor | None: _ExtractValType: TypeAlias = ( None | PySymType - | OpaqueType + | OpaqueBase | _AnyScriptObjectType | BackwardState | list["_ExtractValType"] @@ -656,7 +686,7 @@ def extract_val(val: _ExtractValType, include_real: bool = False) -> _ExtractVal return snapshot_fake(val, include_real=include_real) elif isinstance(val, py_sym_types): return val - elif isinstance(val, _AnyScriptObject): + elif isinstance(val, (_AnyScriptObject, OpaqueBase)): return val elif isinstance(val, BackwardState): return val @@ -904,6 +934,15 @@ def wrap_with_proxy( elif isinstance(e, _AnyScriptObject) or is_opaque_value(e): if not isinstance(proxy, Proxy): raise AssertionError(f"Expected Proxy, got {type(proxy)}") + # Non-hoisted opaque value types should be baked as constants + # in the graph, not tracked as proxy references. This matches + # dynamo's behavior where non-hoisted values are not graph inputs. + if ( + is_opaque_value_type(type(e)) # pyrefly: ignore[bad-argument-type] + and not should_hoist(type(e)) + ): + set_meta(proxy, e) + return set_proxy_slot(e, tracer, proxy) set_meta(proxy, e) elif isinstance(e, (tuple, list)): @@ -1004,13 +1043,13 @@ def fetch_object_proxy( @overload def fetch_object_proxy( - tracer: _ProxyTracer, t: OpaqueType + tracer: _ProxyTracer, t: OpaqueBase ) -> _OpaqueObjectProxyType | PySymType: ... def fetch_object_proxy( tracer: _ProxyTracer, - t: Tensor | _AnyScriptObjectType | PySymType | OpaqueType, + t: Tensor | _AnyScriptObjectType | PySymType | OpaqueBase, ) -> object: return get_proxy_slot(t, tracer, t) @@ -1290,7 +1329,13 @@ def tensor_numel_in_limit(t: Tensor) -> bool: else: constant = None - track_tensor_tree(out, proxy_out, constant=constant, tracer=tracer) + track_tensor_tree( + out, + proxy_out, + # pyrefly: ignore[bad-argument-type] + constant=constant, + tracer=tracer, + ) _maybe_record_pointwise_barrier(func, proxy_mode) return out @@ -1332,35 +1377,42 @@ class _SympyExprTrackerValue: value: PySymType +def _init_proxy_trackers(tracer: PythonKeyTracer | _GraphAppendingTracerEx) -> None: + """Initialize the tracker dictionaries shared by PythonKeyTracer and _GraphAppendingTracerEx.""" + tracer.tensor_tracker = WeakTensorKeyDictionary() + tracer.symnode_tracker = _SymNodeDict() + tracer.script_object_tracker = WeakIdKeyDictionary(dict=None, ref_type=_WeakHashRef) + tracer.opaque_tracker = WeakIdKeyDictionary() + tracer._opaque_real_obj_proxy = {} + tracer.sympy_expr_tracker = {} + # Stores the torch function that was called during tracing + tracer.torch_fn_metadata = None + # Stores the counts for every torch function called. This is to help + # distinguish between different calls to the same torch function. + tracer.torch_fn_counts = {} + tracer.enable_thunkify = False + + class PythonKeyTracer(Tracer): # ScriptObject uses _WeakHashRef because the same C++ IValue can produce # different Python wrapper objects, so Python id() won't match. script_object_tracker: MutableMapping[torch.ScriptObject, Proxy] # FakeScriptObject/OpaqueBase uses WeakIdRef because distinct objects that # are value-equal (e.g. primal vs tangent opaques) must be tracked separately. - opaque_tracker: MutableMapping[FakeScriptObject | OpaqueBase | OpaqueType, Proxy] + opaque_tracker: MutableMapping[FakeScriptObject | OpaqueBase, Proxy] + # Maps id(real_obj) -> proxy for opaque FSOs, so that multiple FSO wrappers + # of the same real object (e.g. primal vs tangent) resolve to one proxy. + _opaque_real_obj_proxy: dict[int, Proxy] symnode_tracker: _SymNodeDict sympy_expr_tracker: dict[sympy.Symbol, _SympyExprTrackerValue] tensor_tracker: MutableMapping[Tensor, _ProxyTensor] + torch_fn_metadata: OpOverload | None torch_fn_counts: dict[OpOverload, int] enable_thunkify: bool = False def __init__(self) -> None: super().__init__(autowrap_modules=()) # type: ignore[arg-type] - self.tensor_tracker = WeakTensorKeyDictionary() - self.symnode_tracker = _SymNodeDict() - self.script_object_tracker = WeakIdKeyDictionary( - dict=None, ref_type=_WeakHashRef - ) - self.opaque_tracker = WeakIdKeyDictionary() - self.sympy_expr_tracker = {} - - # Stores the torch function that was called during tracing - self.torch_fn_metadata = None - # Stores the counts for every torch function called. This is to help - # distinguish between different calls to the same torch function. - self.torch_fn_counts = {} - self.enable_thunkify = False + _init_proxy_trackers(self) # In general, we don't want to make modules leaves. In principle, users of # this tracer might want to override this in order to turn a couple specific @@ -1394,8 +1446,75 @@ def create_arg(self, a: object) -> fx.node.Node: if a.node.constant is None: raise AssertionError("a.node.constant should not be None") return a.node.constant + + # Try reconstructing untracked opaque reference types from existing + # graph inputs (e.g. derive a DeviceMesh submesh from its root mesh). + if isinstance(a, (FakeScriptObject, OpaqueBase)): + node = self._try_reconstruct_opaque(a) + if node is not None: + return node + return super().create_arg(a) # type: ignore[return-value] + def _try_reconstruct_opaque( + self, a: FakeScriptObject | OpaqueBase + ) -> fx.node.Node | None: + """Try to reconstruct an opaque object from existing graph inputs. + + When make_fx encounters an untracked opaque reference type (e.g. a + DeviceMesh submesh captured by a backward closure), this method checks + if the type has a registered reconstruct_fn that can derive the object + from inputs already in the graph. Returns an FX Node on success, + None on failure (falls back to get_attr constant). + """ + real_obj: OpaqueBase = a.real_obj if isinstance(a, FakeScriptObject) else a + + if not is_opaque_reference_type(type(real_obj)): + return None + + reconstruct_fn = get_reconstruct_fn(type(real_obj)) + if reconstruct_fn is None: + return None + + def get_tracked_proxy(obj: OpaqueBase) -> Proxy | None: + proxy = self._opaque_real_obj_proxy.get(id(obj)) + if proxy is not None: + return proxy + proxy = get_proxy_slot(obj, self, None) + if proxy is not None: + return proxy + # The object may be identity-different but equal to a tracked + # FSO's real_obj. This happens because maybe_to_fake_obj creates + # a new DeviceMesh wrapper, so the FSO's real_obj is a distinct + # object from the one held by submesh._root_mesh. Equality + # comparison is needed, not identity. + for tracked_obj, p in self.opaque_tracker.items(): + if not isinstance(tracked_obj, FakeScriptObject): + continue + if tracked_obj.real_obj == obj: + return p + return None + + result = reconstruct_fn(real_obj, get_tracked_proxy, self) + if result is None: + return None + + # The reconstruct_fn dispatches through a custom op with a Proxy + # argument, which goes through Proxy.__torch_function__. That path + # creates a graph node but does NOT populate meta["val"] (unlike the + # ProxyTorchDispatchMode.__torch_dispatch__ path which calls + # track_tensor_tree). Set it here so downstream consumers (e.g. the + # min-cut partitioner) can classify the node correctly. + if "val" not in result.node.meta: + set_meta(result, a) + + # Also register for the *input* object so dedup works for re-encounters. + set_proxy_slot(a, self, result) + if id(real_obj) not in self._opaque_real_obj_proxy: + self._opaque_real_obj_proxy[id(real_obj)] = result + + return result.node + @overload def unwrap_proxy(self, e: Tensor) -> Proxy | Tensor: ... @@ -1547,9 +1666,9 @@ def context_manager_fn() -> Generator[TorchFunctionMode | None, None, None]: @torch._disable_dynamo def dispatch_trace( - root: Module | Callable, + root: Module | Callable[..., Any], tracer: Tracer, - concrete_args: tuple[Any, ...] | None = None, + concrete_args: tuple[object, ...] | None = None, ) -> GraphModule: graph = tracer.trace(root, concrete_args) # type: ignore[arg-type] @@ -1640,21 +1759,22 @@ def get_sym_proxy_slot(t: PySymType) -> Proxy: # TODO: Make downstream users of this work with OperatorBase -ORIGINAL_ATEN: object | None = None +ORIGINAL_ATEN: contextvars.ContextVar[object | None] = contextvars.ContextVar( + "ORIGINAL_ATEN", default=None +) @contextmanager def set_original_aten_op( func: OpOverload | torch._ops.HigherOrderOperator, ) -> Generator[None, None, None]: - global ORIGINAL_ATEN - if ORIGINAL_ATEN is None and fx_traceback.has_preserved_node_meta(): - ORIGINAL_ATEN = func + if ORIGINAL_ATEN.get() is None and fx_traceback.has_preserved_node_meta(): + token = ORIGINAL_ATEN.set(func) fx_traceback.current_meta["original_aten"] = func try: yield finally: - ORIGINAL_ATEN = None + ORIGINAL_ATEN.reset(token) fx_traceback.current_meta["original_aten"] = None else: yield @@ -1696,7 +1816,7 @@ def __init__(self, tracer: _ProxyTracer) -> None: def __torch_function__( self, - func: OpOverload | Callable, + func: OpOverload | Callable[..., Any], types: tuple[torch._C._TensorMeta, ...], args: tuple[object, ...] = (), kwargs: dict[str, object] | None = None, @@ -1763,7 +1883,7 @@ def enable_tracing(self) -> bool: def __init__( self, tracer: _ProxyTracer, - tracing_mode: str, + tracing_mode: _TracingMode, pre_dispatch: bool = False, _allow_fake_constant: bool = False, _error_on_data_dependent_ops: bool = True, @@ -1908,8 +2028,11 @@ def _compute_proxy( class _GraphAppendingTracerEx(fx.proxy.GraphAppendingTracer): script_object_tracker: MutableMapping[torch.ScriptObject, Proxy] - opaque_tracker: MutableMapping[FakeScriptObject | OpaqueBase | OpaqueType, Proxy] - symnode_tracker: MutableMapping[PySymType, _PySymProxyType] + opaque_tracker: MutableMapping[FakeScriptObject | OpaqueBase, Proxy] + # Maps id(real_obj) -> proxy for opaque FSOs, so that multiple FSO wrappers + # of the same real object (e.g. primal vs tangent) resolve to one proxy. + _opaque_real_obj_proxy: dict[int, Proxy] + symnode_tracker: _SymNodeDict tensor_tracker: MutableMapping[Tensor, _ProxyTensor] sympy_expr_tracker: dict[sympy.Symbol, _SympyExprTrackerValue] torch_fn_metadata: OpOverload | None @@ -1918,18 +2041,7 @@ class _GraphAppendingTracerEx(fx.proxy.GraphAppendingTracer): def __init__(self, graph: fx.graph.Graph) -> None: super().__init__(graph) - self.symnode_tracker = weakref.WeakKeyDictionary() - self.tensor_tracker = WeakTensorKeyDictionary() - self.sympy_expr_tracker = {} - self.script_object_tracker = WeakIdKeyDictionary( - dict=None, ref_type=_WeakHashRef - ) - self.opaque_tracker = WeakIdKeyDictionary() - # Stores the torch function that was called during tracing - self.torch_fn_metadata = None - # Stores the counts for every torch function called. This is to help - # distinguish between different calls to the same torch function. - self.torch_fn_counts = {} + _init_proxy_trackers(self) # TODO: I'm not sure what the point of this class is; you can just @@ -1939,7 +2051,7 @@ def __init__( self, module: fx.GraphModule, new_graph: fx.Graph, - decomposition_table: Mapping[OpOverload, Callable] | None = None, + decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None = None, **kwargs: object, ) -> None: super().__init__(module, **kwargs) # type: ignore[arg-type] @@ -2006,7 +2118,7 @@ def __init__( self, module: fx.GraphModule, should_decompose: Callable[[fx.Node], bool], - decomposition_table: Mapping[OpOverload, Callable], + decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None, **kwargs: object, ) -> None: """ @@ -2021,7 +2133,7 @@ def __init__( def recursive_wrap( gm: fx.GraphModule, should_decompose: Callable[[fx.Node], bool], - decomposition_table: Mapping[OpOverload, Callable], + decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None, **kwargs: object, ) -> _SelectiveDecomposeInterpreter: """ @@ -2052,7 +2164,7 @@ def recursive_wrap( gm, should_decompose, decomposition_table, **kwargs ) - def run_node(self, n): + def run_node(self, n: fx.Node) -> Any: if self.should_decompose(n): with decompose(self.decomposition_table): result = super().run_node(n) @@ -2063,9 +2175,9 @@ def run_node(self, n): def selective_decompose( joint_gm: fx.GraphModule, - *args, - decomposition, - should_decompose, + *args: object, + decomposition: Mapping[OpOverload, Callable[..., Any]] | None, + should_decompose: Callable[..., bool], trace_joint_graph: bool, ) -> fx.GraphModule: """Retrace a joint graph module and selectively apply decomposition.""" @@ -2074,13 +2186,13 @@ def selective_decompose( # the arg name, primals and tangents, are important. # make_fx keeps the name in the traced graph and partitioner later relies # on the name to partition joint graph correctly. - def wrap_fn(primals: list[Any], tangents: list[Any]): + def wrap_fn(primals: list[Any], tangents: list[Any]) -> Any: return _SelectiveDecomposeInterpreter.recursive_wrap( joint_gm, should_decompose, decomposition ).run(*args) else: - def wrap_fn(*args): + def wrap_fn(*args: Any) -> Any: return _SelectiveDecomposeInterpreter.recursive_wrap( joint_gm, should_decompose, decomposition ).run(*args) @@ -2285,7 +2397,7 @@ def getattr( return self.attr_proxy_map[attr_val] def trace( # type: ignore[override] - self, root: Module | Callable, concrete_args: dict[str, object] | None + self, root: Module | Callable[..., Any], concrete_args: dict[str, object] | None ) -> fx.Graph: res = super().trace(root, concrete_args) @@ -2356,7 +2468,7 @@ def _delete_proxy_attr(obj: Module, target: str) -> bool: def call_module( self, m: Module, - forward: Callable, + forward: Callable[..., Any], args: tuple[object, ...], kwargs: dict[str, object], ) -> None: @@ -2383,7 +2495,7 @@ def call_module( "This might be because the module was not properly registered " "as a submodule, which is not good practice. We will trace " "through the module without recording stack information.", - str(m), + m, ) return forward(*args, **kwargs) @@ -2415,12 +2527,14 @@ def create_node(self, *args: object, **kwargs: object) -> fx.node.Node: # torch_fn if ( node.op == "call_function" - and self.torch_fn_metadata is not None + and (torch_fn := self.torch_fn_metadata) is not None and "torch_fn" not in node.meta ): node.meta["torch_fn"] = ( - f"{self.torch_fn_metadata.__name__}_{self.torch_fn_counts[self.torch_fn_metadata]}", - f"{self.torch_fn_metadata.__class__.__name__}.{self.torch_fn_metadata.__name__}", + # pyrefly: ignore[missing-attribute,bad-index] + f"{torch_fn.__name__}_{self.torch_fn_counts[torch_fn]}", + # pyrefly: ignore[missing-attribute] + f"{torch_fn.__class__.__name__}.{torch_fn.__name__}", ) return node @@ -2429,8 +2543,8 @@ def create_node(self, *args: object, **kwargs: object) -> fx.node.Node: class _MakefxTracer: def __init__( self, - decomposition_table: Mapping[OpOverload, Callable] | None, - tracing_mode: str, + decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None, + tracing_mode: _TracingMode, _allow_non_fake_inputs: bool, pre_dispatch: bool, record_module_stack: bool, @@ -2443,13 +2557,13 @@ def __init__( ) -> None: # Configurations that are used to initialize the context managers and their states. # Should not modify them during tracing. - self.decomposition_table: dict[OpOverload, Callable] = dict( + self.decomposition_table: dict[OpOverload, Callable[..., Any]] = dict( decomposition_table or {} ) self.decomposition_table.setdefault( torch.ops.aten.sym_numel.default, torch._decomp.decompositions.sym_numel ) - self.tracing_mode: str = tracing_mode + self.tracing_mode: _TracingMode = tracing_mode self._allow_non_fake_inputs: bool = _allow_non_fake_inputs self.pre_dispatch: bool = pre_dispatch self.record_module_stack: bool = record_module_stack @@ -2461,13 +2575,13 @@ def __init__( # Remember to specify how to initialize it from user inputs and from parent tracer whenever # adding new modes in _MakefxTracer. self.fake_tensor_mode: FakeTensorMode | None = None - self.proxy_mode: nullcontext | ProxyTorchDispatchMode = nullcontext() - self.proxy_function_mode: nullcontext | PreDispatchTorchFunctionMode = ( + self.proxy_mode: nullcontext[None] | ProxyTorchDispatchMode = nullcontext() + self.proxy_function_mode: nullcontext[None] | PreDispatchTorchFunctionMode = ( nullcontext() ) self.fx_tracer: PythonKeyTracer | None = None - self.python_dispatcher_mode: nullcontext | Any = nullcontext() - self.torch_fn_metadata_mode: nullcontext | TorchFunctionMetadataMode = ( + self.python_dispatcher_mode: nullcontext[None] | Any = nullcontext() + self.torch_fn_metadata_mode: nullcontext[None] | TorchFunctionMetadataMode = ( nullcontext() ) self.record_stack_traces = record_stack_traces @@ -2488,11 +2602,11 @@ def _checkpoint_modes(self) -> list[Any]: def _restore_modes( self, prev_fake_tensor_mode: FakeTensorMode | None, - prev_proxy_mode: nullcontext | ProxyTorchDispatchMode, - prev_proxy_function_mode: nullcontext | PreDispatchTorchFunctionMode, + prev_proxy_mode: nullcontext[None] | ProxyTorchDispatchMode, + prev_proxy_function_mode: nullcontext[None] | PreDispatchTorchFunctionMode, prev_fx_tracer: PythonKeyTracer | None, - prev_python_dispatcher_mode: nullcontext | Any, - prev_torch_fn_metadata_mode: nullcontext | TorchFunctionMetadataMode, + prev_python_dispatcher_mode: nullcontext[None] | Any, + prev_torch_fn_metadata_mode: nullcontext[None] | TorchFunctionMetadataMode, ) -> None: self.fake_tensor_mode = prev_fake_tensor_mode self.proxy_mode = prev_proxy_mode @@ -2503,7 +2617,7 @@ def _restore_modes( @contextmanager def _init_modes_from_inputs( - self, f: Callable, args: tuple[object, ...] + self, f: Callable[..., Any], args: tuple[object, ...] ) -> Generator[None, None, None]: prev_modes = self._checkpoint_modes() try: @@ -2618,7 +2732,57 @@ def _create_sub_fx_tracer(parent_tracer: _ProxyTracer) -> PythonKeyTracer: finally: self._restore_modes(*prev_modes) - def _trace_inner(self, f: Callable, *args: object) -> GraphModule: + def _convert_args_to_fake(self, args: T) -> T: + if self.tracing_mode == "real": + return args + + arg_count = 0 + + def inner_wrap_fake(x: object) -> object: + nonlocal arg_count + # TODO: it would be nice to line these up with the names + # FX will choose for the placeholders, but we don't + # actually know what the names will be at this point yet + # NB: the Source here is actually meaningless + from torch._dynamo.source import ConstantSource + + if self.fake_tensor_mode is None: + raise AssertionError("fake_tensor_mode should not be None") + source = ConstantSource(f"input{arg_count}") + if isinstance(x, Tensor): + arg_count += 1 + return self.fake_tensor_mode.from_tensor(x, source=source) + # NB: don't match on bools + elif type(x) is int and self.tracing_mode == "symbolic": + if self.fake_tensor_mode.shape_env is None: + raise AssertionError( + "shape_env should be set if tracing with 'symbolic'" + ) + return self.fake_tensor_mode.shape_env.create_symintnode( + self.fake_tensor_mode.shape_env.create_symbol( + x, source, positive=None + ), + hint=x, + source=source, + ) + elif isinstance(x, torch.ScriptObject) or is_opaque_value(x): + if is_opaque_value_type( + type(x) # pyrefly: ignore[bad-argument-type] + ): + return x + return torch._library.fake_class_registry.maybe_to_fake_obj( + self.fake_tensor_mode, x + ) + + if isinstance(x, FakeScriptObject): + raise AssertionError( + f"ScriptObject {x} has been fakified. Cannot wrap_fake it again." + ) + return x + + return pytree.tree_map(inner_wrap_fake, args) + + def _trace_inner(self, f: Callable[..., Any], *args: object) -> GraphModule: # TODO: We need to explicitly import torch._dynamo before calling dispatch_trace, # because dispatch_trace will introduce the lazy import of torch._dynamo, # and some contexts set before calling dispatch_trace will cause problems with the import of torch._dynamo, @@ -2628,70 +2792,16 @@ def _trace_inner(self, f: Callable, *args: object) -> GraphModule: phs = pytree.tree_map(lambda _: torch.fx._symbolic_trace.PH, args) - def _wrap_fake(args: T) -> T: - arg_count = 0 - - def inner_wrap_fake(x: object) -> object: - nonlocal arg_count - # TODO: it would be nice to line these up with the names - # FX will choose for the placeholders, but we don't - # actually know what the names will be at this point yet - # NB: the Source here is actually meaningless - from torch._dynamo.source import ConstantSource - - if self.fake_tensor_mode is None: - raise AssertionError("fake_tensor_mode should not be None") - source = ConstantSource(f"input{arg_count}") - if isinstance(x, Tensor): - arg_count += 1 - return self.fake_tensor_mode.from_tensor(x, source=source) - # NB: don't match on bools - elif type(x) is int and self.tracing_mode == "symbolic": - if self.fake_tensor_mode.shape_env is None: - raise AssertionError( - "shape_env should be set if tracing with 'symbolic'" - ) - return self.fake_tensor_mode.shape_env.create_symintnode( - self.fake_tensor_mode.shape_env.create_symbol( - x, source, positive=None - ), - hint=x, - source=source, - ) - elif isinstance(x, torch.ScriptObject) or is_opaque_value(x): - if is_opaque_value_type( - type(x) # pyrefly: ignore[bad-argument-type] - ): - return x - return torch._library.fake_class_registry.maybe_to_fake_obj( - self.fake_tensor_mode, x - ) + args = self._convert_args_to_fake(args) - if isinstance(x, FakeScriptObject): - raise AssertionError( - f"ScriptObject {x} has been fakified. Cannot wrap_fake it again." - ) - return x - - wrap_fn_map = { - "real": lambda x: x, - "fake": inner_wrap_fake, - "symbolic": inner_wrap_fake, - } - return pytree.tree_map(wrap_fn_map[self.tracing_mode], args) - - def _wrap_func(f: Callable[_P, R], phs: Sequence[PHBase]) -> Callable[_P, R]: - if ( - not hasattr(inspect.unwrap(f), "__code__") - or inspect.unwrap(f).__code__.co_flags & inspect.CO_VARARGS - ): - # FX doesn't support varargs, so we gotta fake up a wrapper - # TODO: Would be nice to fix this at the source... - return fake_signature(f, len(phs)) - return f - - args = _wrap_fake(args) - func = _wrap_func(f, phs) + # FX doesn't support varargs, so we gotta fake up a wrapper + # TODO: Would be nice to fix this at the source... + func: Callable[..., Any] = f + if ( + not hasattr(inspect.unwrap(f), "__code__") + or inspect.unwrap(f).__code__.co_flags & inspect.CO_VARARGS + ): + func = fake_signature(f, len(phs)) # We disable the autocast cache as the autocast cache causes type conversions on parameters to # check a cache, which introduces untracked tensors into the graph # @@ -2752,18 +2862,21 @@ def _wrap_func(f: Callable[_P, R], phs: Sequence[PHBase]) -> Callable[_P, R]: t.shape_env = self.fake_tensor_mode.shape_env # type: ignore[assignment] return t - def trace(self, f: Callable, *args: object) -> fx.GraphModule: + def trace(self, f: Callable[..., Any], *args: object) -> fx.GraphModule: with self._init_modes_from_inputs(f, args): return self._trace_inner(f, *args) def is_hop_subgraph_tracer(self) -> bool: return self.parent_tracer is not None - def trace_subgraph(self, f: Callable, *args: object) -> GraphModule: - # Create a new tracer based on parent's config - sub_tracer = _MakefxTracer( - self.decomposition_table, - "real", + def _make_sub_tracer( + self, + decomp_table: Mapping[OpOverload, Callable[..., Any]] | None = None, + tracing_mode: _TracingMode = "real", + ) -> _MakefxTracer: + return _MakefxTracer( + decomp_table if decomp_table is not None else self.decomposition_table, + tracing_mode, self._allow_non_fake_inputs, self.pre_dispatch, self.record_module_stack, @@ -2771,25 +2884,21 @@ def trace_subgraph(self, f: Callable, *args: object) -> GraphModule: self._error_on_data_dependent_ops, parent_tracer=self, ) + + def trace_subgraph(self, f: Callable[..., Any], *args: object) -> GraphModule: + sub_tracer = self._make_sub_tracer() with sub_tracer._init_modes_from_parent(self): return sub_tracer._trace_inner(f, *args) def trace_subgraph_custom_decomp( - self, f: Callable, decomp_table: Mapping[OpOverload, Callable], *args + self, + f: Callable[..., Any], + decomp_table: Mapping[OpOverload, Callable[..., Any]], + *args: object, ) -> GraphModule: if not isinstance(decomp_table, Mapping): raise AssertionError(f"Expected Mapping, got {type(decomp_table)}") - # Create a new tracer based on parent's config, but use a different decomposition table - sub_tracer = _MakefxTracer( - decomp_table, - "real", - self._allow_non_fake_inputs, - self.pre_dispatch, - self.record_module_stack, - self._allow_fake_constant, - self._error_on_data_dependent_ops, - parent_tracer=self, - ) + sub_tracer = self._make_sub_tracer(decomp_table=decomp_table) with sub_tracer._init_modes_from_parent(self): return sub_tracer._trace_inner(f, *args) @@ -2809,9 +2918,9 @@ def _set_make_fx_tracer(tracer: _MakefxTracer) -> Generator[None, None, None]: def make_fx( - f: Callable, - decomposition_table: Mapping[OpOverload, Callable] | None = None, - tracing_mode: str = "real", + f: Callable[..., Any], + decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None = None, + tracing_mode: _TracingMode = "real", _allow_non_fake_inputs: bool = False, *, pre_dispatch: bool = False, @@ -2917,7 +3026,8 @@ def maybe_handle_decomp( ) -> object: from torch._inductor.compiler_bisector import CompilerBisector - if op in CURRENT_DECOMPOSITION_TABLE: + decomp_table = CURRENT_DECOMPOSITION_TABLE.get({}) + if op in decomp_table: if CompilerBisector.disable_subsystem( "aot_eager_decomp_partition", "decomposition", lambda: repr(op) ): @@ -2925,7 +3035,7 @@ def maybe_handle_decomp( with proxy_mode: proxy_mode.decomp_layers += 1 - out = CURRENT_DECOMPOSITION_TABLE[op](*args, **kwargs) + out = decomp_table[op](*args, **kwargs) proxy_mode.decomp_layers -= 1 return out @@ -2933,11 +3043,11 @@ def maybe_handle_decomp( def get_isolated_graphmodule( - func: Callable, + func: Callable[..., Any], args: tuple[object, ...], kwargs: dict[str, object], - tracing_mode: str = "real", - decomposition_table: Mapping[OpOverload, Callable] | None = None, + tracing_mode: _TracingMode = "real", + decomposition_table: Mapping[OpOverload, Callable[..., Any]] | None = None, ) -> GraphModule: """A helper function used to get the GraphModule for the given func. diff --git a/torch/fx/experimental/recording.py b/torch/fx/experimental/recording.py index 28ff27ab3f2d1..0d401df3d18ad 100644 --- a/torch/fx/experimental/recording.py +++ b/torch/fx/experimental/recording.py @@ -1,16 +1,26 @@ -# mypy: allow-untyped-defs +from __future__ import annotations + import functools import inspect import itertools import logging -from collections.abc import Callable from dataclasses import dataclass -from typing import Any +from typing import Any, ParamSpec, TYPE_CHECKING, TypeVar + + +_P = ParamSpec("_P") +_R = TypeVar("_R") import torch import torch.utils._pytree as pytree +if TYPE_CHECKING: + from collections.abc import Callable + + from torch.fx.experimental.symbolic_shapes import ShapeEnv, TrackedFake + + log = logging.getLogger(__name__) trace_shape_events_log = torch._logging.getArtifactLogger( __name__, "trace_shape_events" @@ -81,21 +91,21 @@ @dataclass class ShapeEnvEvent: # ShapeEnv method. - f: Callable + f: Callable[..., Any] # Arguments and keyword arguments called with. - args: list[Any] | None = None + args: list[object] | None = None kwargs: dict[str, Any] | None = None # List of tracked_fakes at the time the method was called. - tracked_fakes: list[Any] | None = None + tracked_fakes: list[TrackedFake] | None = None # Name of the captured event. # Used for special handling of particular methods. name: str | None = None # Replay itself, but using shape_env as self. - def run(self, shape_env=None) -> Any: + def run(self, shape_env: ShapeEnv | None = None) -> Any: from torch.fx.experimental.symbolic_shapes import ( is_symbolic, ShapeEnv, @@ -147,7 +157,7 @@ def maybe_convert_node(x: Any) -> Any: return name_to_node[x.name] # Replaces the value of an specific argument by the result of fn. - def replacearg(index: int, key: str, fn: Callable): + def replacearg(index: int, key: str, fn: Callable[..., Any]) -> None: if index < len(args): args[index] = fn(args[index]) if key in kwargs: @@ -196,7 +206,9 @@ def is_defer_runtime_assert(self) -> bool: # 2. SymInt, SymFloat, or SymBool arguments # If we find more than one object of any of the above types, we # also check that the ShapeEnv instance is the same for all of them. -def _extract_shape_env_and_assert_equal(args, kwargs): +def _extract_shape_env_and_assert_equal( + args: tuple[object, ...] | list[object], kwargs: dict[str, object] +) -> ShapeEnv | None: from torch.fx.experimental.symbolic_shapes import is_symbolic, ShapeEnv, SymTypes def assert_equal(old: ShapeEnv | None, new: ShapeEnv) -> ShapeEnv: @@ -241,8 +253,8 @@ def assert_equal(old: ShapeEnv | None, new: ShapeEnv) -> ShapeEnv: # - ShapeEnv.guard_or_defer_runtime_assert def record_shapeenv_event( *, save_tracked_fakes: bool = False, name: str | None = None -) -> Callable: - def decorator(fn: Callable) -> Callable: +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: + def decorator(fn: Callable[_P, _R]) -> Callable[_P, _R]: if not callable(fn): raise AssertionError(f"Expected callable, got {type(fn)}") args = inspect.getfullargspec(fn).args @@ -256,7 +268,7 @@ def decorator(fn: Callable) -> Callable: name = fn.__name__ @functools.wraps(fn) - def wrapper(*args, **kwargs): + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: from torch.fx.experimental.symbolic_shapes import ShapeEnv if not isinstance(args[0], ShapeEnv): @@ -269,7 +281,7 @@ def wrapper(*args, **kwargs): ) NEST += 1 - def retlog(r): + def retlog(r: _R) -> _R: trace_shape_events_log.debug("%s-> %s", " " * (NEST - 1), r) return r @@ -346,7 +358,7 @@ def retlog(r): # It assumes the first event is the constructor call. # # fn: transforms an old FX node into one corresponding to the newly created ShapeEnv. -def replay_shape_env_events(events): +def replay_shape_env_events(events: list[ShapeEnvEvent]) -> ShapeEnv: from torch.fx.experimental.symbolic_shapes import ShapeEnv constructor_event = events[0] @@ -394,7 +406,7 @@ def dim(self) -> int: return len(self.tensor_size) @staticmethod - def from_fake(fake) -> "FakeTensorMeta": + def from_fake(fake: torch.Tensor) -> FakeTensorMeta: return FakeTensorMeta( fake.size(), fake.stride(), fake.storage_offset(), fake.is_nested ) @@ -446,7 +458,12 @@ def from_fake(fake) -> "FakeTensorMeta": # Checks whether the state of two ShapeEnv are equal w.r.t. the guards # returned by ShapeEnv.produce_guards. -def shape_env_check_state_equal(env1, env2, non_state_variable_names, map_value): +def shape_env_check_state_equal( + env1: ShapeEnv, + env2: ShapeEnv, + non_state_variable_names: tuple[str, ...], + map_value: Callable[[str, object], object], +) -> None: # Collect and remove variables that don't necessarily represent the state # of a ShapeEnv. Note: we copy the dictionary so that we don't modify the # instance itself. @@ -476,7 +493,7 @@ def value_to_str(value: Any) -> str: # Here, we allow the value of each field to be mapped, so that we appropriately # compare the two values. def compare_vars( - map_value: Callable[[str, Any], Any], + map_value: Callable[[str, object], object], ) -> list[tuple[str, str, str]]: env1_set, env2_set = set(env1_vars), set(env2_vars) diff --git a/torch/fx/experimental/schema_type_annotation.py b/torch/fx/experimental/schema_type_annotation.py index b507c17822116..9377f0bafdfcd 100644 --- a/torch/fx/experimental/schema_type_annotation.py +++ b/torch/fx/experimental/schema_type_annotation.py @@ -1,15 +1,21 @@ # mypy: allow-untyped-defs +from __future__ import annotations + import inspect -from typing import Any +from typing import Any, TYPE_CHECKING import torch import torch.fx from torch._jit_internal import boolean_dispatched from torch.fx import Transformer -from torch.fx.node import Argument, Target from torch.fx.operator_schemas import _torchscript_type_to_python_type +if TYPE_CHECKING: + from torch.fx.graph_module import GraphModule + from torch.fx.node import Argument, Target + + class AnnotateTypesWithSchema(Transformer): """ Use Python function signatures to annotate types for `Nodes` within an FX graph. @@ -31,7 +37,7 @@ class AnnotateTypesWithSchema(Transformer): def __init__( self, - module: torch.nn.Module, + module: GraphModule, annotate_functionals: bool = True, annotate_modules: bool = True, annotate_get_attrs: bool = True, diff --git a/torch/fx/experimental/sym_node.py b/torch/fx/experimental/sym_node.py index 21b124b7f1539..a175c92df0526 100644 --- a/torch/fx/experimental/sym_node.py +++ b/torch/fx/experimental/sym_node.py @@ -1,5 +1,3 @@ -# mypy: allow-untyped-defs - from __future__ import annotations @@ -23,7 +21,7 @@ import operator import sys from functools import lru_cache, update_wrapper -from typing import TYPE_CHECKING +from typing import Any, overload, TYPE_CHECKING import torch import torch._logging.structured as structured @@ -43,6 +41,11 @@ if TYPE_CHECKING: + from collections.abc import Callable + from typing import Self + + import sympy + from torch.fx.experimental.symbolic_shapes import ShapeEnv log = logging.getLogger(__name__) @@ -62,7 +65,15 @@ from torch.types import py_sym_types as SymTypes -def _to_symtype(t): +@overload +def _to_symtype(t: type[bool]) -> type[SymBool]: ... +@overload +def _to_symtype(t: type[int]) -> type[SymInt]: ... +@overload +def _to_symtype(t: type[float]) -> type[SymFloat]: ... +@overload +def _to_symtype(t: type) -> type: ... +def _to_symtype(t: type) -> type: if t is bool: return SymBool if t is int: @@ -92,14 +103,14 @@ class SymNode: def __init__( self, - expr, - shape_env, - pytype, + expr: object, + shape_env: ShapeEnv | None, + pytype: type, hint: HintType | object, - constant=None, - fx_node=None, - optimized_summation=False, - ): + constant: int | float | bool | None = None, + fx_node: object = None, + optimized_summation: bool = False, + ) -> None: self._expr = expr self.shape_env = shape_env self.pytype = pytype @@ -134,7 +145,7 @@ def __init__( # potential refinements to unbacked symints this got harder to keep # in sync, so we've deleted it for now.) - def compute_hint(): + def compute_hint() -> HintType | SymInt | SymFloat | SymBool: from torch.fx.experimental.symbolic_shapes import has_free_unbacked_symbols # This occasionally gets exercised by, e.g., @@ -144,6 +155,8 @@ def compute_hint(): # expensive. if has_free_unbacked_symbols(self.expr): return None + if self.shape_env is None: + raise RuntimeError("shape_env is required to compute hint") hint = self.shape_env._maybe_evaluate_static(self.expr, compute_hint=True) if hint is not None: hint = self.pytype(hint) if not isinstance(hint, SymTypes) else hint @@ -197,9 +210,14 @@ def _value_hash(self) -> int: return hash((self._expr, self.pytype, self._hint, self.constant, self.fx_node)) @property - def expr(self): - if isinstance(self._expr, int) or self._expr.is_number: + def expr(self) -> sympy.Basic: + if ( + isinstance(self._expr, int) + or self._expr.is_number # pyrefly: ignore[missing-attribute] + ): return self._expr + if self.shape_env is None: + raise AssertionError("shape_env is required to access expr") ver = self.shape_env._replacements_version_counter if ver == 0: return self._expr @@ -211,20 +229,20 @@ def expr(self): return result @property - def hint(self): + def hint(self) -> HintType | SymInt | SymFloat | SymBool: return self._hint - def has_hint(self): + def has_hint(self) -> bool: return self._hint is not None - def maybe_as_int(self): + def maybe_as_int(self) -> int | None: if self.expr.is_number: return int(self.expr) else: return None # NB: This does conversions, not sure if this is good or not - def maybe_as_float(self): + def maybe_as_float(self) -> float | None: import sympy if isinstance(self.expr, sympy.Float): @@ -232,7 +250,7 @@ def maybe_as_float(self): else: return None - def maybe_as_bool(self): + def maybe_as_bool(self) -> bool | None: import sympy if self.expr is sympy.true: @@ -242,16 +260,16 @@ def maybe_as_bool(self): else: return None - def is_int(self): + def is_int(self) -> bool: return self.pytype is int - def is_float(self): + def is_float(self) -> bool: return self.pytype is float - def is_bool(self): + def is_bool(self) -> bool: return self.pytype is bool - def is_nested_int(self): + def is_nested_int(self) -> bool: # Unbacked SymInts cannot be nested int today return ( self._hint is not None @@ -259,7 +277,7 @@ def is_nested_int(self): and self._hint.node.is_nested_int() ) - def wrap_int(self, num): + def wrap_int(self, num: int) -> SymNode: if type(num) is not int: raise AssertionError(f"Expected int, got {type(num)}") import sympy @@ -268,7 +286,7 @@ def wrap_int(self, num): sympy.Integer(num), self.shape_env, int, num, constant=num, fx_node=num ) - def wrap_float(self, num): + def wrap_float(self, num: float) -> SymNode: if type(num) is not float: raise AssertionError(f"Expected float, got {type(num)}") import sympy @@ -277,7 +295,7 @@ def wrap_float(self, num): sympy.Float(num), self.shape_env, float, num, constant=num, fx_node=num ) - def wrap_bool(self, num): + def wrap_bool(self, num: bool) -> SymNode: if type(num) is not bool: raise AssertionError(f"Expected bool, got {type(num)}") import sympy @@ -291,16 +309,16 @@ def wrap_bool(self, num): fx_node=num, ) - def clone(self): + def clone(self) -> SymNode: return self - def str(self): + def str(self) -> builtins.str: return f"{self.expr}" - def __str__(self): + def __str__(self) -> builtins.str: return self.str() - def __repr__(self): + def __repr__(self) -> builtins.str: rep = [ f"SymNode({self._expr}, shape_env={self.shape_env}, pytype={self.pytype}", ] @@ -324,70 +342,70 @@ def abs(self) -> SymNode: def pos(self) -> SymNode: return self._pos() # type: ignore[attr-defined] - def round(self, ndigits=None) -> SymNode: + def round(self, ndigits: int | None = None) -> SymNode: return self._round(ndigits) # type: ignore[attr-defined] def trunc(self) -> SymNode: return self._trunc() # type: ignore[attr-defined] - def add(self, other) -> SymNode: + def add(self, other: SymNode) -> SymNode: return self._add(other) # type: ignore[attr-defined] - def sub(self, other) -> SymNode: + def sub(self, other: SymNode) -> SymNode: return self._sub(other) # type: ignore[attr-defined] - def mul(self, other) -> SymNode: + def mul(self, other: SymNode) -> SymNode: return self._mul(other) # type: ignore[attr-defined] - def mod(self, other) -> SymNode: + def mod(self, other: SymNode) -> SymNode: return self._mod(other) # type: ignore[attr-defined] - def float_pow(self, other) -> SymNode: + def float_pow(self, other: SymNode) -> SymNode: return self._float_pow(other) # type: ignore[attr-defined] - def pow_by_natural(self, other) -> SymNode: + def pow_by_natural(self, other: SymNode) -> SymNode: return self._pow_by_natural(other) # type: ignore[attr-defined] - def and_(self, other) -> SymNode: + def and_(self, other: SymNode) -> SymNode: return self._and_(other) # type: ignore[attr-defined] - def or_(self, other) -> SymNode: + def or_(self, other: SymNode) -> SymNode: return self._or_(other) # type: ignore[attr-defined] - def float_truediv(self, other) -> SymNode: + def float_truediv(self, other: SymNode) -> SymNode: return self._float_truediv(other) # type: ignore[attr-defined] - def int_truediv(self, other) -> SymNode: + def int_truediv(self, other: SymNode) -> SymNode: return self._int_truediv(other) # type: ignore[attr-defined] - def int_floordiv(self, other) -> SymNode: + def int_floordiv(self, other: SymNode) -> SymNode: return self._int_floordiv(other) # type: ignore[attr-defined] - def lshift(self, other) -> SymNode: + def lshift(self, other: SymNode) -> SymNode: return self._lshift(other) # type: ignore[attr-defined] - def rshift(self, other) -> SymNode: + def rshift(self, other: SymNode) -> SymNode: return self._rshift(other) # type: ignore[attr-defined] def sym_not(self) -> SymNode: # noqa: F811 return self._sym_not() # type: ignore[attr-defined] - def eq(self, other) -> SymNode: + def eq(self, other: SymNode) -> SymNode: return self._eq(other) # type: ignore[attr-defined] - def ne(self, other) -> SymNode: + def ne(self, other: SymNode) -> SymNode: return self._ne(other) # type: ignore[attr-defined] - def gt(self, other) -> SymNode: + def gt(self, other: SymNode) -> SymNode: return self._gt(other) # type: ignore[attr-defined] - def lt(self, other) -> SymNode: + def lt(self, other: SymNode) -> SymNode: return self._lt(other) # type: ignore[attr-defined] - def le(self, other) -> SymNode: + def le(self, other: SymNode) -> SymNode: return self._le(other) # type: ignore[attr-defined] - def ge(self, other) -> SymNode: + def ge(self, other: SymNode) -> SymNode: return self._ge(other) # type: ignore[attr-defined] def floor(self) -> SymNode: @@ -408,74 +426,86 @@ def ceil(self) -> SymNode: def neg(self) -> SymNode: return self._neg() # type: ignore[attr-defined] - def sym_min(self, other) -> SymNode: # noqa: F811 + def sym_min(self, other: SymNode) -> SymNode: # noqa: F811 return self._sym_min(other) # type: ignore[attr-defined] - def sym_max(self, other) -> SymNode: # noqa: F811 + def sym_max(self, other: SymNode) -> SymNode: # noqa: F811 return self._sym_max(other) # type: ignore[attr-defined] - def sym_ite(self, then_val, else_val) -> SymNode: + def sym_ite(self, then_val: SymNode, else_val: SymNode) -> SymNode: return self._sym_ite(then_val, else_val) # type: ignore[attr-defined] - def is_contiguous(self, sizes, strides) -> SymNode: + def is_contiguous(self, sizes: list[SymNode], strides: list[SymNode]) -> SymNode: return self._is_contiguous(sizes, strides) # type: ignore[attr-defined] - def is_channels_last_contiguous_2d(self, sizes, strides) -> SymNode: + def is_channels_last_contiguous_2d( + self, sizes: list[SymNode], strides: list[SymNode] + ) -> SymNode: return self._is_channels_last_contiguous_2d(sizes, strides) # type: ignore[attr-defined] - def is_channels_last_contiguous_3d(self, sizes, strides) -> SymNode: + def is_channels_last_contiguous_3d( + self, sizes: list[SymNode], strides: list[SymNode] + ) -> SymNode: return self._is_channels_last_contiguous_3d(sizes, strides) # type: ignore[attr-defined] - def is_channels_last_strides_2d(self, sizes, strides) -> SymNode: + def is_channels_last_strides_2d( + self, sizes: list[SymNode], strides: list[SymNode] + ) -> SymNode: return self._is_channels_last_strides_2d(sizes, strides) # type: ignore[attr-defined] - def is_channels_last_strides_3d(self, sizes, strides) -> SymNode: + def is_channels_last_strides_3d( + self, sizes: list[SymNode], strides: list[SymNode] + ) -> SymNode: return self._is_channels_last_strides_3d(sizes, strides) # type: ignore[attr-defined] - def is_non_overlapping_and_dense_indicator(self, sizes, strides) -> SymNode: + def is_non_overlapping_and_dense_indicator( + self, sizes: list[SymNode], strides: list[SymNode] + ) -> SymNode: return self._is_non_overlapping_and_dense_indicator(sizes, strides) # type: ignore[attr-defined] # Make C++ happy - def sym_or(self, other): + def sym_or(self, other: SymNode) -> SymNode: return self.or_(other) - def sym_and(self, other): + def sym_and(self, other: SymNode) -> SymNode: return self.and_(other) # Integer bitwise ops - def bitwise_and(self, other): + def bitwise_and(self, other: SymNode) -> SymNode: return self._bitwise_and(other) # type: ignore[attr-defined] - def bitwise_or(self, other): + def bitwise_or(self, other: SymNode) -> SymNode: return self._bitwise_or(other) # type: ignore[attr-defined] - def bitwise_xor(self, other): + def bitwise_xor(self, other: SymNode) -> SymNode: return self._bitwise_xor(other) # type: ignore[attr-defined] # There is no int_truediv available from C++ - def truediv(self, other): + def truediv(self, other: SymNode) -> SymNode: return self.float_truediv(other) - def floordiv(self, other) -> SymNode: + def floordiv(self, other: SymNode) -> SymNode: return self.int_floordiv(other) # We didn't bind integer pow in C++ - def pow(self, other): + def pow(self, other: SymNode) -> SymNode: return self.float_pow(other) - def is_non_overlapping_and_dense(self, sizes, strides): + def is_non_overlapping_and_dense( + self, sizes: list[SymNode], strides: list[SymNode] + ) -> SymNode: return self.is_non_overlapping_and_dense_indicator(sizes, strides).eq( to_node(self, 1) ) # type: ignore[attr-defined] - def int_(self): + def int_(self) -> int: return self.guard_int("", 0) # NB: uses Python backtrace # This one is currently done by hand, but if we add other variadic # functions consider factoring it out to be metaprogrammed too. Note that # some load bearing logic is directly in torch.sym_sum - def sym_sum(self, args) -> SymNode: + def sym_sum(self, args: list[SymNode]) -> SymNode: import sympy # Inner impl @@ -503,8 +533,10 @@ def sym_sum(self, args) -> SymNode: break size_hints.append(a.hint) else: - out_hint = sum(size_hints) + out_hint = sum(size_hints) # pyrefly: ignore[no-matching-overload] + if self.shape_env is None: + raise RuntimeError("shape_env is required for sym_sum") fx_node, _ = self.shape_env._create_fx_call_function( torch.sym_sum, (tuple(a.fx_node for a in args),) ) @@ -512,11 +544,13 @@ def sym_sum(self, args) -> SymNode: # NB: Only for integers! return SymNode(out, self.shape_env, int, out_hint, fx_node=fx_node) - def evaluate(self, size_oblivious=False): + def evaluate(self, size_oblivious: bool = False) -> bool | int | float: + if self.shape_env is None: + raise RuntimeError("shape_env is required to evaluate") return self.shape_env.evaluate_sym_node(self, size_oblivious) # You can manually trigger a guard with this function - def guard_int(self, file, line): + def guard_int(self, file: builtins.str, line: int) -> int: # TODO: use the file/line for some useful diagnostic on why a # guard occurred r = self.evaluate() @@ -526,7 +560,7 @@ def guard_int(self, file, line): log.warning("Failed to convert to int: %s", r) raise - def guard_float(self, file, line): + def guard_float(self, file: builtins.str, line: int) -> float: # TODO: use the file/line for some useful diagnostic on why a # guard occurred r = self.evaluate() @@ -536,7 +570,7 @@ def guard_float(self, file, line): log.warning("Failed to convert to float: %s", r) raise - def guard_bool(self, file, line): + def guard_bool(self, file: builtins.str, line: int) -> bool: # TODO: use the file/line for some useful diagnostic on why a # guard occurred r = self.evaluate() @@ -546,9 +580,11 @@ def guard_bool(self, file, line): log.warning("Failed to convert to bool: %s", r) raise - def expect_true(self, file, line): + def expect_true(self, file: builtins.str, line: int) -> bool: from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols + if self.shape_env is None: + raise RuntimeError("shape_env is required for expect_true") if ( self.has_hint() and not free_unbacked_symbols(self.expr) @@ -564,14 +600,14 @@ def expect_true(self, file, line): self.expr, f"{file}:{line}", fx_node=self.fx_node ) - def statically_known_true(self, file, line): + def statically_known_true(self, file: builtins.str, line: int) -> bool: from torch.fx.experimental.symbolic_shapes import statically_known_true if not self.is_bool(): raise AssertionError("Expected bool type") return statically_known_true(SymBool(self)) - def guard_size_oblivious(self, file, line): + def guard_size_oblivious(self, file: builtins.str, line: int) -> bool: """ Like guard_bool, but if we encounter unbacked symbols, if those symbols are size-like, we will treat them as >= 2 for the purposes of the analysis. @@ -591,35 +627,35 @@ def guard_size_oblivious(self, file, line): log.warning("Failed to convert to bool: %s", r) raise - def guard_or_false(self, file, line): + def guard_or_false(self, file: builtins.str, line: int) -> bool: from torch.fx.experimental.symbolic_shapes import guard_or_false if not self.is_bool(): raise AssertionError("Expected bool type") return guard_or_false(SymBool(self)) - def guard_or_true(self, file, line): + def guard_or_true(self, file: builtins.str, line: int) -> bool: from torch.fx.experimental.symbolic_shapes import guard_or_true if not self.is_bool(): raise AssertionError("Expected bool type") return guard_or_true(SymBool(self)) - def bool_(self): + def bool_(self) -> bool: return self.guard_bool("", 0) - def is_symbolic(self): + def is_symbolic(self) -> bool: return True - def nested_int(self): + def nested_int(self) -> None: return None - def is_constant(self): + def is_constant(self) -> bool: return False class _DynamicScalar: - def __new__(cls, *args): + def __new__(cls, *args: object) -> Self: if cls is _DynamicScalar: raise TypeError("_DynamicScalar is an abstract base class, use DynamicInt.") return super().__new__(cls, *args) @@ -637,21 +673,44 @@ class DynamicInt(_DynamicScalar, int): fn(x) # compiles x as a dynamic integer input; returns f(4) """ - def __new__(cls, val): + def __new__(cls, val: int) -> Self: if not isinstance(val, int): raise AssertionError(f"Expected int, got {type(val)}") obj = super().__new__(cls, int(val)) return obj - def __repr__(self): + def __repr__(self) -> str: return f"DynamicInt({self.real})" - def __floordiv__(self, other): # // was casting to int without these overrides? + def __floordiv__( + self, other: int + ) -> DynamicInt: # // was casting to int without these overrides? return DynamicInt(self.real // other) - def __rfloordiv__(self, other): + def __rfloordiv__(self, other: int) -> DynamicInt: return DynamicInt(other // self.real) + def __pow__(self, other, modulo=None): + if modulo is not None: + result = pow(self.real, other, modulo) + else: + result = self.real**other + # Only create DynamicInt if result is int, otherwise return plain value + # (e.g., negative exponent produces float) + if isinstance(result, int): + return DynamicInt(result) + return result + + def __rpow__(self, other, modulo=None): + if modulo is not None: + result = pow(other, self.real, modulo) + else: + result = other**self.real + # Only create DynamicInt if result is int, otherwise return plain value + if isinstance(result, int): + return DynamicInt(result) + return result + # TODO: this probably needs the sizes-strides eval functions METHOD_TO_OPERATOR = { @@ -706,8 +765,8 @@ def __rfloordiv__(self, other): # Adding math ops: sqrt, cos, sin, ... -def _get_sym_node_fn(name): - def fn(self): +def _get_sym_node_fn(name: str) -> Callable[[SymNode], SymNode]: + def fn(self: SymNode) -> SymNode: return getattr(self, f"_sym_{name}")() return fn @@ -732,6 +791,7 @@ def fn(self): setattr(SymNode, sym_name, _get_sym_node_fn(name)) METHOD_TO_OPERATOR[sym_name] = getattr(torch, priv_sym_name) unary_magic_methods.add(sym_name) + # pyrefly: ignore [unresolvable-dunder-all] __all__.append(sym_name) @@ -785,25 +845,25 @@ def fn(self): # Methods that have a `__foo__` as well as `__rfoo__` -def _sympy_float_truediv(a, b): +def _sympy_float_truediv(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import FloatTrueDiv return FloatTrueDiv(a, b) -def _sympy_int_truediv(a, b): +def _sympy_int_truediv(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import IntTrueDiv return IntTrueDiv(a, b) -def _sympy_floordiv(a, b): +def _sympy_floordiv(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import FloorDiv return FloorDiv(a, b) -def _sympy_mod(a, b): +def _sympy_mod(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import Mod, PythonMod if a.is_nonnegative and b.is_nonnegative: @@ -812,43 +872,45 @@ def _sympy_mod(a, b): return PythonMod(a, b) -def _sympy_pow_by_natural(a, b): +def _sympy_pow_by_natural(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import PowByNatural return PowByNatural(a, b) -def _sympy_float_pow(a, b): +def _sympy_float_pow(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import FloatPow return FloatPow(a, b) -def _sympy_and(a, b): +def _sympy_and(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: import sympy return sympy.And(a, b) -def _sympy_or(a, b): +def _sympy_or(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: import sympy return sympy.Or(a, b) -def _sympy_lshift(a, b): +def _sympy_lshift(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import LShift return LShift(a, b) -def _sympy_rshift(a, b): +def _sympy_rshift(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import RShift return RShift(a, b) -def _binary_search_insert_arg(ordered_args, new_arg): +def _binary_search_insert_arg( + ordered_args: list[sympy.Basic], new_arg: sympy.Basic +) -> list[sympy.Basic] | None: """ If new_arg is found in ordered_args None is returned, else the new ordered_args with new_arg inserted @@ -883,8 +945,11 @@ def _binary_search_insert_arg(ordered_args, new_arg): def _optimized_add( - lhs, rhs, lhs_is_optimized_summation=False, rhs_is_optimized_summation=False -): + lhs: sympy.Basic, + rhs: sympy.Basic, + lhs_is_optimized_summation: bool = False, + rhs_is_optimized_summation: bool = False, +) -> tuple[bool, sympy.Basic]: """ Custom optimization for Add used to optimize incremental binary summations of certain properties. The idea is when we know the expression is a summation of unique symbols all we need to know is the correct order of symbols, @@ -898,7 +963,7 @@ def _optimized_add( import sympy from sympy.core.basic import _args_sortkey as sortkey - def make_optimized(ordered_args): + def make_optimized(ordered_args: list[sympy.Basic]) -> tuple[bool, sympy.Basic]: if ordered_args is None: raise AssertionError("ordered_args is None") # Use _from_args directly to bypass _exec_constructor_postprocessors @@ -950,19 +1015,19 @@ def make_optimized(ordered_args): return (_is_symbols_binary_summation(result), result) -def _bitwise_and(a, b): +def _bitwise_and(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import BitwiseFn_bitwise_and return BitwiseFn_bitwise_and(a, b) -def _bitwise_or(a, b): +def _bitwise_or(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import BitwiseFn_bitwise_or return BitwiseFn_bitwise_or(a, b) -def _bitwise_xor(a, b): +def _bitwise_xor(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import BitwiseFn_bitwise_xor return BitwiseFn_bitwise_xor(a, b) @@ -988,7 +1053,7 @@ def _bitwise_xor(a, b): } -def _floor_ceil_helper(a, fn): +def _floor_ceil_helper(a: sympy.Basic, fn: Callable[..., sympy.Basic]) -> sympy.Basic: import sympy if isinstance(a, sympy.Mul): @@ -1006,7 +1071,7 @@ def _floor_ceil_helper(a, fn): return fn(a) -def _sympy_floor(a): +def _sympy_floor(a: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import FloorToInt return FloorToInt(a) @@ -1014,67 +1079,67 @@ def _sympy_floor(a): # NB: this is Python trunc semantics which returns an int. Do NOT use this to # represent torch.trunc (which is float to float) -def _sympy_trunc(a): +def _sympy_trunc(a: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import TruncToInt return TruncToInt(a) -def _sympy_ceil(a): +def _sympy_ceil(a: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import CeilToInt return CeilToInt(a) -def _sympy_eq(a, b): +def _sympy_eq(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: import sympy return sympy.Eq(a, b) -def _sympy_ne(a, b): +def _sympy_ne(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: import sympy return sympy.Ne(a, b) -def _sympy_gt(a, b): +def _sympy_gt(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: import sympy return sympy.Gt(a, b) -def _sympy_lt(a, b): +def _sympy_lt(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: import sympy return sympy.Lt(a, b) -def _sympy_le(a, b): +def _sympy_le(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: import sympy return sympy.Le(a, b) -def _sympy_ge(a, b): +def _sympy_ge(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: import sympy return sympy.Ge(a, b) -def _sympy_min(a, b): +def _sympy_min(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import Min return Min(a, b) -def _sympy_max(a, b): +def _sympy_max(a: sympy.Basic, b: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import Max return Max(a, b) -def _sympy_ite(a, t, f): +def _sympy_ite(a: sympy.Basic, t: sympy.Basic, f: sympy.Basic) -> sympy.Basic: import sympy return sympy.Piecewise((t, a), (f, True)) @@ -1083,8 +1148,8 @@ def _sympy_ite(a, t, f): current_module = sys.modules[__name__] -def _get_sym_math_fn(name): - def fn(a): +def _get_sym_math_fn(name: str) -> Callable[[sympy.Basic], sympy.Basic]: + def fn(a: sympy.Basic) -> sympy.Basic: import torch.utils._sympy.functions return getattr(torch.utils._sympy.functions, f"OpaqueUnaryFn_{name}")(a) @@ -1101,13 +1166,15 @@ def fn(a): del fn, name, priv_sympy_name # type: ignore[possibly-undefined] -def _sympy_abs(a): +def _sympy_abs(a: sympy.Basic) -> sympy.Basic: import sympy return sympy.Abs(a) -def _sympy_round(number, ndigits=None): +def _sympy_round( + number: sympy.Basic, ndigits: sympy.Basic | None = None +) -> sympy.Basic: from torch.utils._sympy.functions import RoundDecimal, RoundToInt if ndigits is None: @@ -1116,7 +1183,7 @@ def _sympy_round(number, ndigits=None): return RoundDecimal(number, ndigits) -def _sympy_sym_float(a): +def _sympy_sym_float(a: sympy.Basic) -> sympy.Basic: from torch.utils._sympy.functions import ToFloat # NB: Cannot use a * 1.0 here, because 0 * 1.0 is 0 which incorrectly @@ -1124,7 +1191,7 @@ def _sympy_sym_float(a): return ToFloat(a) -def _sympy_is_integer(a): +def _sympy_is_integer(a: sympy.Basic) -> sympy.Basic: import sympy from torch.utils._sympy.functions import ToFloat @@ -1163,12 +1230,16 @@ def _sympy_is_integer(a): del name, sym_name, math_op_names, current_module # type: ignore[possibly-undefined] -def sympy_is_contiguous(sizes, strides): +def sympy_is_contiguous( + sizes: list[sympy.Basic], strides: list[sympy.Basic] +) -> sympy.Basic: dim = len(sizes) return sympy_is_contiguous_generic(sizes, strides, list(range(dim - 1, -1, -1))) -def sympy_is_contiguous_generic(sizes, strides, dim_order): +def sympy_is_contiguous_generic( + sizes: list[sympy.Basic], strides: list[sympy.Basic], dim_order: list[int] +) -> sympy.Basic: import sympy dim = len(sizes) @@ -1192,15 +1263,21 @@ def sympy_is_contiguous_generic(sizes, strides, dim_order): # happens you will need to refactor this -def sympy_is_channels_last_contiguous_2d(sizes, strides): +def sympy_is_channels_last_contiguous_2d( + sizes: list[sympy.Basic], strides: list[sympy.Basic] +) -> sympy.Basic: return sympy_is_contiguous_generic(sizes, strides, [1, 3, 2, 0]) -def sympy_is_channels_last_contiguous_3d(sizes, strides): +def sympy_is_channels_last_contiguous_3d( + sizes: list[sympy.Basic], strides: list[sympy.Basic] +) -> sympy.Basic: return sympy_is_contiguous_generic(sizes, strides, [1, 4, 3, 2, 0]) -def sympy_is_channels_last_strides_generic(sizes, strides, dim_order): +def sympy_is_channels_last_strides_generic( + sizes: list[sympy.Basic], strides: list[sympy.Basic], dim_order: list[int] +) -> sympy.Basic: import sympy from torch.utils._sympy.functions import Max @@ -1240,15 +1317,21 @@ def sympy_is_channels_last_strides_generic(sizes, strides, dim_order): return r -def sympy_is_channels_last_strides_2d(sizes, strides): +def sympy_is_channels_last_strides_2d( + sizes: list[sympy.Basic], strides: list[sympy.Basic] +) -> sympy.Basic: return sympy_is_channels_last_strides_generic(sizes, strides, [1, 3, 2, 0]) -def sympy_is_channels_last_strides_3d(sizes, strides): +def sympy_is_channels_last_strides_3d( + sizes: list[sympy.Basic], strides: list[sympy.Basic] +) -> sympy.Basic: return sympy_is_channels_last_strides_generic(sizes, strides, [1, 4, 3, 2, 0]) -def _sympy_is_non_overlapping_and_dense_indicator(sizes, strides): +def _sympy_is_non_overlapping_and_dense_indicator( + sizes: list[sympy.Basic], strides: list[sympy.Basic] +) -> sympy.Basic: from torch.utils._sympy.functions import IsNonOverlappingAndDenseIndicator return IsNonOverlappingAndDenseIndicator(*sizes, *strides) @@ -1266,7 +1349,7 @@ def _sympy_is_non_overlapping_and_dense_indicator(sizes, strides): } -def to_node(self, num): +def to_node(self: SymNode, num: object) -> SymNode: if isinstance(num, SymTypes): return num.node elif type(num) is bool: @@ -1281,7 +1364,7 @@ def to_node(self, num): return NotImplemented -def wrap_node(x): +def wrap_node(x: SymNode) -> SymInt | SymFloat | SymBool | int | float | bool: # TODO: let C++ also take advantage of this if isinstance(x, SymNode) and x.constant is not None: return x.constant @@ -1295,11 +1378,11 @@ def wrap_node(x): raise AssertionError(f"unrecognized return type {x}") -def method_to_operator(method): +def method_to_operator(method: str) -> Callable[..., object]: return METHOD_TO_OPERATOR[method] -def _make_node_magic(method, func): +def _make_node_magic(method: str, func: Callable[..., sympy.Basic]) -> None: func = lru_cache(256)(func) if method in magic_methods_on_operator_with_trailing_underscore: @@ -1324,9 +1407,9 @@ def uninteresting_files() -> set[str]: | {""} ) - def capture_provenance(fn): + def capture_provenance(fn: Callable[..., SymNode]) -> Callable[..., SymNode]: @functools.wraps(fn) - def wrapper(self, other=None): + def wrapper(self: SymNode, other: SymNode | None = None) -> SymNode: if other is None: result = fn(self) else: @@ -1337,7 +1420,7 @@ def wrapper(self, other=None): else: arguments = [self] - def get_id(sym_node) -> int | None: + def get_id(sym_node: SymNode) -> int | None: # We don't want to return an ID if the input is a constant import sympy @@ -1371,7 +1454,7 @@ def get_id(sym_node) -> int | None: return wrapper @capture_provenance - def binary_magic_impl(self, other): + def binary_magic_impl(self: SymNode, other: SymNode) -> SymNode: from torch.fx.experimental.proxy_tensor import ( get_proxy_mode, handle_sym_dispatch, @@ -1397,6 +1480,8 @@ def binary_magic_impl(self, other): # Special handling for mod that requires access to the value # ranges shape_env = self.shape_env + if shape_env is None: + raise AssertionError("shape_env is required for mod") if ( self.expr.is_nonnegative or shape_env.bound_sympy(self.expr).lower >= 0 @@ -1480,6 +1565,8 @@ def binary_magic_impl(self, other): # Create a FX node that corresponds to the operation being applied to # this node. + if self.shape_env is None: + raise RuntimeError("shape_env is required for binary op") fx_node, _ = self.shape_env._create_fx_call_function( op, (self.fx_node, other.fx_node) ) @@ -1495,7 +1582,7 @@ def binary_magic_impl(self, other): return result @capture_provenance - def unary_magic_impl(self): + def unary_magic_impl(self: SymNode) -> SymNode: from torch.fx.experimental.proxy_tensor import ( get_proxy_mode, handle_sym_dispatch, @@ -1506,6 +1593,8 @@ def unary_magic_impl(self): return to_node(self, handle_sym_dispatch(op, (wrap_node(self),), {})) # TODO: consider constant prop here expr = self.expr + if self.shape_env is None: + raise RuntimeError("shape_env is required for unary op") if method == "floor" or method == "ceiling": expr = self.shape_env._simplify_floor_div(expr) @@ -1535,13 +1624,20 @@ def unary_magic_impl(self): setattr(SymNode, f"_{method_attr}", unary_magic_impl) elif method == "sym_ite": - def sym_ite_impl(pred_node, then_node, else_node): + def sym_ite_impl( + pred_node: SymNode, then_node: SymNode, else_node: SymNode + ) -> SymNode: from torch.fx.experimental.proxy_tensor import ( get_proxy_mode, handle_sym_dispatch, ) - out_hint = then_node.hint if pred_node.hint else else_node.hint + if pred_node.hint is None: + out_hint = None + elif pred_node.hint: + out_hint = then_node.hint + else: + out_hint = else_node.hint if get_proxy_mode(): return to_node( pred_node, @@ -1568,6 +1664,8 @@ def sym_ite_impl(pred_node, then_node, else_node): ) raise + if pred_node.shape_env is None: + raise RuntimeError("shape_env is required for sym_ite") fx_node, _ = pred_node.shape_env._create_fx_call_function( sym_ite, (pred_node.fx_node, then_node.fx_node, else_node.fx_node) ) @@ -1578,7 +1676,7 @@ def sym_ite_impl(pred_node, then_node, else_node): setattr(SymNode, f"_{method_attr}", sym_ite_impl) elif method == "round": - def round_impl(self, ndigits=None): + def round_impl(self: SymNode, ndigits: int | None = None) -> SymNode: from torch.fx.experimental.proxy_tensor import ( get_proxy_mode, handle_sym_dispatch, @@ -1604,7 +1702,9 @@ def round_impl(self, ndigits=None): out_hint = None if self.hint is not None: - out_hint = op(self.hint, ndigits) + out_hint = op( # pyrefly: ignore[no-matching-overload] + self.hint, ndigits + ) # Internally, None is used as sentinel to indicate that a something is not a node on an FX graph. At the # same time, there is no way to wrap a plain None into an FX node. Thus, there is no way to pass None here @@ -1616,6 +1716,8 @@ def round_impl(self, ndigits=None): args = [self.fx_node] if ndigits is not None: args.append(ndigits) + if self.shape_env is None: + raise RuntimeError("shape_env is required for round") fx_node, _ = self.shape_env._create_fx_call_function(op, tuple(args)) return SymNode(out, self.shape_env, pytype, out_hint, fx_node=fx_node) @@ -1624,10 +1726,12 @@ def round_impl(self, ndigits=None): setattr(SymNode, f"_{method_attr}", binary_magic_impl) -def _make_node_sizes_strides(method, func): +def _make_node_sizes_strides(method: str, func: Callable[..., sympy.Basic]) -> None: # NB: don't LRU cache, lots of arguments - def sizes_strides_impl(self, sizes, strides): + def sizes_strides_impl( + self: SymNode, sizes: list[SymNode], strides: list[SymNode] + ) -> SymNode: from torch.fx.experimental.proxy_tensor import ( get_proxy_mode, handle_sym_dispatch, @@ -1680,7 +1784,9 @@ def sizes_strides_impl(self, sizes, strides): # TODO: This is technically hotpath, but in the ideal end state # guards on this will resolve at a higher level so you never # spend time in this code - def sizes_strides_user(sizes, strides): + def sizes_strides_user( + sizes: list[object], strides: list[object] + ) -> SymInt | SymFloat | SymBool | int | float | bool: import sympy from torch.fx.experimental.symbolic_shapes import ( @@ -1696,7 +1802,10 @@ def sizes_strides_user(sizes, strides): ) ) if method == "is_non_overlapping_and_dense_indicator": - return eval_is_non_overlapping_and_dense(sizes, strides) + return eval_is_non_overlapping_and_dense( + sizes, # pyrefly: ignore[bad-argument-type] + strides, # pyrefly: ignore[bad-argument-type] + ) else: # TODO: this is an awful implementation return bool( @@ -1718,7 +1827,7 @@ def sizes_strides_user(sizes, strides): _make_node_sizes_strides(method, func) -def _make_user_magic(method, user_type): +def _make_user_magic(method: str, user_type: type) -> None: # User magic takes care of wrapping the other operand into a node, # so that our internal logic can assume everything is nodes if method in magic_methods_on_operator_with_trailing_underscore: @@ -1726,7 +1835,9 @@ def _make_user_magic(method, user_type): else: method_attr = method - def get_constant(x: SymInt | int | SymFloat | float | SymBool | bool): + def get_constant( + x: SymInt | int | SymFloat | float | SymBool | bool, + ) -> int | float | bool: if isinstance(x, (int, float, bool)): return x if isinstance(x, SymInt): @@ -1735,7 +1846,7 @@ def get_constant(x: SymInt | int | SymFloat | float | SymBool | bool): return x.node.guard_bool("", 0) raise AssertionError("expect to be called with constant SymBools") - def is_constant(x): + def is_constant(x: SymInt | int | SymFloat | float | SymBool | bool) -> bool: if isinstance(x, (int, float, bool)): return True if isinstance(x, (SymInt, SymFloat, SymBool)): @@ -1770,7 +1881,7 @@ def is_constant(x): if method in bool_becomes_int_magic_methods: - def promote(x): + def promote(x: object) -> Any: """Implements True+True=2, which works in python but not sympy""" if isinstance(x, SymBool): return SymInt(x.node.wrap_int(int(x))) @@ -1778,10 +1889,10 @@ def promote(x): else: - def promote(x): + def promote(x: object) -> Any: return x - def promote2(self, other): + def promote2(self: object, other: object) -> tuple[Any, Any]: # TODO: Remove eq and other relations from this list. # CPython has fancy implementations for these to get as much precision # as possible instead of just promoting to float64 and praying, so we @@ -1823,13 +1934,13 @@ def promote2(self, other): # Alternatively, we could also rewrap into constant Symbool (i.e. by # implementing wrap_bool in ConstantSymNodeImpl), but we're not doing that # today for no particular reason. - def unary_magic_impl(self): + def unary_magic_impl(self: object) -> Any: self = promote(self) if is_constant(self): return (method_to_operator(method))(get_constant(self)) return wrap_node(getattr(self.node, method_attr)()) - def binary_magic_impl(self, other): + def binary_magic_impl(self: object, other: object) -> Any: if not isinstance(other, (int, float, bool, SymInt, SymFloat, SymBool)): return NotImplemented sym_node_log.debug("MAGIC %s %s %s", method, self, other) @@ -1846,7 +1957,7 @@ def binary_magic_impl(self, other): ret = wrap_node(getattr(self.node, method_attr)(other_node)) return get_constant(ret) if is_constant(ret) else ret - def rbinary_magic_impl(self, other): + def rbinary_magic_impl(self: object, other: object) -> Any: if not isinstance(other, (int, float, bool, SymInt, SymFloat, SymBool)): return NotImplemented self = promote(self) @@ -1862,7 +1973,7 @@ def rbinary_magic_impl(self, other): ret = wrap_node(getattr(other_node, method_attr)(self.node)) return get_constant(ret) if is_constant(ret) else ret - def setattrs(user_type, attr, symnode_impl): + def setattrs(user_type: type, attr: str, symnode_impl: object) -> None: """ Registers the SymNode magic method on SymInt/Float/Bool, and optionally registers a corresponding wrapped method on DynamicInt. @@ -1872,8 +1983,10 @@ def setattrs(user_type, attr, symnode_impl): setattr(user_type, attr, symnode_impl) # DynamicInt impl - def dynamic_int_impl(*args): - args = [x.real if isinstance(x, DynamicInt) else x for x in args] + def dynamic_int_impl(*args: object) -> Any: + args = [ # pyrefly: ignore[bad-assignment] + x.real if isinstance(x, DynamicInt) else x for x in args + ] out = getattr(int, attr)(*args) if isinstance(out, int) and not isinstance(out, bool): return DynamicInt(out) @@ -1889,7 +2002,9 @@ def dynamic_int_impl(*args): setattrs(user_type, method, update_wrapper(unary_magic_impl, orig)) elif method == "sym_ite": - def sym_ite_magic_impl(pred, then_val, else_val): + def sym_ite_magic_impl( + pred: SymBool, then_val: object, else_val: object + ) -> Any: pred_node = pred.node then_node = to_node(pred_node, then_val) else_node = to_node(pred_node, else_val) @@ -1904,12 +2019,17 @@ def sym_ite_magic_impl(pred, then_val, else_val): "then_node and else_node must be SymNodes with same pytype" ) ret = wrap_node(getattr(pred.node, method_attr)(then_node, else_node)) - return get_constant(ret) if ret.node.is_constant() else ret + return ( + get_constant(ret) + # pyrefly: ignore[missing-attribute] + if ret.node.is_constant() + else ret + ) setattrs(user_type, f"__{method}__", sym_ite_magic_impl) elif method == "round": - def round_magic_impl(self, ndigits=None): + def round_magic_impl(self: SymFloat, ndigits: int | None = None) -> Any: if is_constant(self): return builtins.round(get_constant(self), ndigits) diff --git a/torch/fx/experimental/symbolic_shapes.py b/torch/fx/experimental/symbolic_shapes.py index 58dd5c93c15e9..d239bae2a5628 100644 --- a/torch/fx/experimental/symbolic_shapes.py +++ b/torch/fx/experimental/symbolic_shapes.py @@ -136,7 +136,9 @@ def guarding_hint_or_throw( """ if isinstance(a, SymNode): if a._hint is not None: - return a._hint + return a._hint # pyrefly: ignore[bad-return] + if a.shape_env is None: + raise AssertionError("shape_env is required for guarding_hint_or_throw") hint = a.shape_env.guarding_hint_or_throw(a.expr) a._hint = hint return hint @@ -943,7 +945,7 @@ def div_by_factor(x: sympy.Expr, factor: int) -> sympy.Expr: return expr -def is_nested_int(s: IntLikeType) -> TypeGuard[SymInt]: +def is_nested_int(s: IntLikeType | FloatLikeType) -> TypeGuard[SymInt]: return isinstance(s, torch.SymInt) and s.node.is_nested_int() @@ -1142,7 +1144,7 @@ def find_symbol_binding_fx_nodes( Returns: A dictionary mapping from sympy Symbols to their binding FX nodes """ - r = {} + r: dict[sympy.Symbol, torch.fx.Node] = {} # NB: Prefer first occurrence of symbol for node in graph.nodes: if (s := is_symbol_binding_fx_node(node)) is not None and s not in r: @@ -1161,7 +1163,7 @@ class Specialization: """ source: TensorPropertySource - check_fn: Callable + check_fn: Callable[[int], bool] # Analogous to ConvertIntSource @@ -1256,7 +1258,7 @@ def expr(s: SymInt | SymFloat | SymBool) -> sympy.Expr: pending = set() r = {} - def match_tensor(a: torch.Tensor, real_tensor: torch.Tensor | None = None): + def match_tensor(a: torch.Tensor, real_tensor: torch.Tensor | None = None) -> None: r.update( go( a.size(), @@ -1733,6 +1735,7 @@ def _advise_is_size(a: SymInt) -> None: isinstance(a, SymInt) and isinstance(a.node, SymNode) and isinstance(a.node.expr, sympy.Symbol) + and a.node.shape_env is not None and a.node.shape_env.is_unbacked_symint(a.node.expr) ): _constrain_range_for_size(a) @@ -1743,6 +1746,7 @@ def _advise_is_bounded(a: SymInt, upper_bound: IntLikeType) -> None: isinstance(a, SymInt) and isinstance(a.node, SymNode) and isinstance(a.node.expr, sympy.Symbol) + and a.node.shape_env is not None and a.node.shape_env.is_unbacked_symint(a.node.expr) and isinstance(upper_bound, int) # TODO: relax ): @@ -2008,7 +2012,7 @@ class StrictMinMaxConstraint(Constraint): for N=0/1 too. """ - vr: ValueRanges + vr: ValueRanges[sympy.Expr] def render(self, source: Source) -> str: """Format the constrain equation""" @@ -2291,7 +2295,7 @@ def __post_init__(self) -> None: # is used. # TODO(voz): Shape env validation @dataclass(frozen=True, slots=True, kw_only=True) -class StatefulSymbolicContext(StatelessSymbolicContext): +class StatefulSymbolicContext(StatelessSymbolicContext[..., Any]): """ Create symbols in ``create_symbolic_sizes_strides_storage_offset`` via a symbolic_context determination as given by a cache of Source:Symbol. A cache hit @@ -2380,7 +2384,8 @@ def _expandsums(args: list[sympy.Expr]) -> tuple[sympy.Expr, bool]: - A boolean indicating whether expansion occurred (True if multiple additive expressions were present or if there was at least one additive and one other expression) """ - adds, other = [], [] + adds: list[sympy.Expr] = [] + other: list[sympy.Expr] = [] for arg in args: if arg.is_Add: adds.append(arg) @@ -2474,7 +2479,7 @@ def safe_expand(r: _SympyT) -> _SympyT: class _SymbolInfo(NamedTuple): k: sympy.Symbol - vr: ValueRanges | None + vr: ValueRanges[sympy.Expr] | None val: sympy.Integer | None is_size_like: bool @@ -2664,8 +2669,16 @@ def cast_symbool_to_symint_guardless( ) +def _eval_is_non_overlapping_and_dense_flat(*args: int) -> int: + # Guard code strings print IsNonOverlappingAndDenseIndicator with flat args + # (s0, s1, ..., stride0, stride1, ...) but eval_is_non_overlapping_and_dense + # expects two sequences (sizes, strides). This wrapper bridges the gap. + dim = len(args) // 2 + return eval_is_non_overlapping_and_dense(list(args[:dim]), list(args[dim:])) + + SYMPY_INTERP = { - "IsNonOverlappingAndDenseIndicator": eval_is_non_overlapping_and_dense, + "IsNonOverlappingAndDenseIndicator": _eval_is_non_overlapping_and_dense_flat, "cast_symbool_to_symint_guardless": cast_symbool_to_symint_guardless, "math": math, "torch": torch, @@ -2944,7 +2957,7 @@ class _CppShapeGuardsHelper(_ShapeGuardsHelper): class LoggingShapeGuardPrinter(ShapeGuardPythonPrinter): - def __init__(self, var_to_sources: Mapping[sympy.Symbol, list[Source]]): + def __init__(self, var_to_sources: Mapping[sympy.Symbol, list[Source]]) -> None: super().__init__(var_to_sources, lambda n: n.name, var_to_sources) @@ -2961,7 +2974,7 @@ def __init__( self, symbol_to_source: dict[sympy.Symbol, list[Source]], source_name_to_debug_name: Mapping[str, str], - ): + ) -> None: super().__init__() self.symbol_to_source = symbol_to_source self.source_name_to_debug_name = source_name_to_debug_name @@ -3202,7 +3215,7 @@ def add_equality(self, source: Source, expr: sympy.Expr) -> None: def _reduce_congruences(self) -> dict[sympy.Symbol, set[sympy.Expr]]: reduced_congruences: dict[sympy.Symbol, set[sympy.Expr]] = {} for s, congruences in self._congruences.items(): - remainder_modulus_pairs = [] + remainder_modulus_pairs: list[tuple[sympy.Integer, sympy.Integer]] = [] congruences_to_check = set() for congruence in congruences: base, divisor = congruence.args @@ -3679,8 +3692,8 @@ def relation_with_digit(expr: str, op: str, digit: int) -> None: self._process_derived_dim_roots(results, name_to_dim) - dims = [] - others = [] + dims: list[str] = [] + others: list[str] = [] # order results by source name results2 = { @@ -3751,7 +3764,7 @@ class ValueRangesSLoc: @contextmanager -def _suppress_guards(shape_env: ShapeEnv) -> Iterator[None]: +def _suppress_guards(shape_env: ShapeEnv) -> Generator[None, None, None]: shape_env._suppress_guards_enter() try: yield @@ -3914,7 +3927,7 @@ def _init( # are conservative: the int MUST fall in the range, but the # range may contain ints which may not actually appear in # practice - self.var_to_range: dict[sympy.Symbol, ValueRanges] = {} + self.var_to_range: dict[sympy.Symbol, ValueRanges[sympy.Expr]] = {} self.var_to_range_sloc: dict[sympy.Symbol, ValueRangesSLoc] = {} self.source_name_to_debug_name: dict[str, str] = {} self.var_to_sources: dict[sympy.Symbol, list[Source]] = {} @@ -3947,7 +3960,7 @@ def _init( self.size_like: set[sympy.Symbol] = set() # Duck-shaping says that if two input tensors have the same size, # they get assigned the same symbolic variable - self.val_to_var: dict[int, sympy.Symbol] = {} + self.val_to_var: dict[IntLikeType | FloatLikeType, sympy.Symbol] = {} self.unbacked_symfloat_counter = 0 self.unbacked_symint_counter = 0 # Similar to guards, but these MUST evaluate to true and can @@ -4038,7 +4051,9 @@ def _init( # 2. list of arguments # This drastically reduces the size of the FX graph, avoiding # duplicated nodes. - self.fx_node_cache: dict[tuple[Callable, tuple[Any, ...]], torch.fx.Node] = {} + self.fx_node_cache: dict[ + tuple[Callable[..., Any], tuple[Any, ...]], torch.fx.Node + ] = {} self.source_to_symbol: dict[str, sympy.Symbol] = {} # Suppose you want to replace an unbacked symbol with another @@ -4135,7 +4150,7 @@ def prefer_deferred_runtime_asserts_over_guards(self) -> bool: @contextmanager def patch_source_specialization( self, source: Source, check_fn: Callable[[sympy.Symbol], sympy.Expr] - ) -> Iterator[None]: + ) -> Generator[None, None, None]: """ Temporarily add symbol-level axioms to the ShapeEnv. This is useful when you want to "fork" and have parallel universes of ShapeEnvs. For example, we use this when doing multi-graph @@ -4265,7 +4280,7 @@ def _last_event_index(self) -> int: return len(self.events) - 1 @contextmanager - def _recording(self) -> Iterator[None]: + def _recording(self) -> Generator[None, None, None]: self.is_recording = True try: yield @@ -4277,7 +4292,9 @@ def _eliminate_unbacked(self, orig_s: sympy.Symbol, new_s: sympy.Expr) -> None: self._set_replacement(orig_s, new_s, "eliminate_unbacked") @record_shapeenv_event() - def set_real_tensor_prop_unbacked_vals(self, k: sympy.Symbol, v: int) -> None: + def set_real_tensor_prop_unbacked_vals( + self, k: sympy.Symbol, v: int | float | torch.Tensor + ) -> None: """Used only when propagate_real_tensors; registers a value for an unbacked symbol, which can be used last resort to resolve hints.""" log.info("set_real_tensor_prop_unbacked_vals %s = %s", k, v) @@ -4403,7 +4420,7 @@ def _ignore_fresh_unbacked_symbols_set(self, b: bool) -> bool: return prev @contextmanager - def ignore_fresh_unbacked_symbols(self) -> Iterator[None]: + def ignore_fresh_unbacked_symbols(self) -> Generator[None, None, None]: """ Indicates that the newly allocated unbacked SymInts are being discarded @@ -4462,8 +4479,8 @@ def _check_translation_validate(self) -> None: @record_shapeenv_event() def _create_fx_call_function( self, - op: Callable, - args: tuple, + op: Callable[..., object], + args: tuple[Any, ...], ) -> tuple[torch.fx.Node | None, bool]: # Cache this tuple in order to avoid duplicated nodes. node_key = (op, args) @@ -4558,7 +4575,7 @@ def suppress_guards(self) -> _GeneratorContextManager[None]: return _suppress_guards(self) @contextmanager - def error_on_new_guards(self) -> Iterator[None]: + def error_on_new_guards(self) -> Generator[None, None, None]: """Context manager that raises _ShapeEnvGuardError if a guard is attempted. Temporarily freezes the ShapeEnv and makes _check_frozen raise @@ -4973,7 +4990,7 @@ def create_symintnode( self, sym: sympy.Expr, *, - hint: int | None, + hint: int | float | bool | torch.SymInt | None, source: Source | None = None, ) -> IntLikeType: """Create a SymInt value from a symbolic expression @@ -5019,7 +5036,7 @@ def create_symfloatnode( self, sym: sympy.Expr, *, - hint: int | float | bool | None, + hint: int | float | bool | torch.SymInt | None, source: Source | None = None, ) -> FloatLikeType: """Create a SymFloat value from a symbolic expression""" @@ -5092,7 +5109,7 @@ def _log_create_unbacked_symbol( self, prefix: str, symbol: sympy.Symbol, - vr: ValueRanges, + vr: ValueRanges[sympy.Expr], source: Source | None = None, sym_node: SymNode | None = None, ) -> None: @@ -5216,7 +5233,7 @@ def create_unspecified_symbol( source: Source, dynamic_dim: DimDynamic = DimDynamic.DUCK, constraint_dim: DimConstraint = None, # NB: includes None - symbolic_context: StatelessSymbolicContext | None = None, + symbolic_context: SymbolicContext | None = None, ) -> sympy.Expr: """ Create a symbol with an unspecified value @@ -5242,13 +5259,13 @@ def create_unspecified_symbol( @record_shapeenv_event() def create_symbol( self, - val: int, + val: IntLikeType | FloatLikeType, source: Source, dynamic_dim: DimDynamic = DimDynamic.DUCK, constraint_dim: DimConstraint = None, # NB: includes None positive: bool | None = True, do_not_specialize_zero_one: bool = False, - symbolic_context: StatelessSymbolicContext | None = None, + symbolic_context: SymbolicContext | None = None, ) -> sympy.Expr: """Create a new symbol which is tracked by this ShapeEnv""" # check if constraint_dim is actually static integer @@ -5262,7 +5279,7 @@ def create_symbol( f"Static shape constraint of {constraint_dim.vr.lower} does not match input size of {val}, " f"for {source.name}" ) - if symbolic_context: + if isinstance(symbolic_context, StatelessSymbolicContext): from torch._dynamo.source import TensorPropertySource if not isinstance(source, TensorPropertySource): @@ -5669,7 +5686,9 @@ def produce_guards_verbose( raise AssertionError(f"len({placeholders}) != len({sources})") Tensorlike = (torch.Tensor, FakeTensorMeta) - def _create_no_constraints_context(t: Tensor) -> StatelessSymbolicContext: + def _create_no_constraints_context( + t: Tensor, + ) -> StatelessSymbolicContext[..., Any]: return StatelessSymbolicContext( # Ignored; only the constraints part is relevant below. dynamic_sizes=[DimDynamic.DYNAMIC] * t.dim(), @@ -5993,7 +6012,7 @@ def track_symfloat(source: Source, val: FloatLikeType) -> None: s = sympy.Float(val) input_guards.append((source, s)) - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] for t, source, context in zip(placeholders, sources, input_contexts): if isinstance(source, str): from torch._dynamo.source import LocalSource @@ -6289,7 +6308,7 @@ def issue_guard(guard: ShapeGuard) -> None: if not sources: raise AssertionError(f"sources must not be empty for symbol {symbol}") - bounds = [] + bounds: list[sympy.Basic] = [] rf = source_ref(sources[0]) verbose_expr = "" if r.lower not in (-sympy.oo, -int_oo): @@ -6680,7 +6699,7 @@ def format_guards(self, verbose: bool = False) -> str: def bound_sympy( self, expr: sympy.Expr, size_oblivious: bool = False - ) -> ValueRanges: + ) -> ValueRanges[sympy.Expr]: """Given a sympy expression, computes a ValueRanges bound for what values it can be""" # TODO: maybe it's guaranteed x in is var_to_range? var_to_range = {x: self.var_to_range.get(x, None) for x in expr.free_symbols} @@ -6843,7 +6862,7 @@ def _maybe_evaluate_static( compute_hint: bool = False, size_oblivious: bool = False, axioms: tuple[SympyBoolean] | None = None, - var_to_range: tuple[tuple[sympy.Symbol, ValueRanges]] | None = None, + var_to_range: tuple[tuple[sympy.Symbol, ValueRanges[sympy.Expr]]] | None = None, ) -> sympy.Basic | None: """ Tries to evaluate expr without introducing guards @@ -6936,7 +6955,7 @@ def replace(self, expr: _SympyT) -> _SympyT: tracks only replacement changes) to cache calls to this method, so depending on other state would cause stale cache results. """ - replacements = {} + replacements: dict[sympy.Basic, sympy.Basic] = {} # pyrefly: ignore [missing-attribute] for s in expr.free_symbols: r = self._find(s) @@ -6972,7 +6991,7 @@ def simplify(self, expr: _SympyT, size_oblivious: bool = False) -> _SympyT: # Simplify max(0/1, x) to x when x >= 0/1. max(1, x) is a commonly introduced # expression when creating contiguous strides. if not size_oblivious: - min_max_replacements = {} + min_max_replacements: dict[sympy.Basic, sympy.Basic] = {} for atom in expr.atoms(Max): # type: ignore[has-type] if len(atom.args) > 2: continue @@ -6988,7 +7007,7 @@ def simplify(self, expr: _SympyT, size_oblivious: bool = False) -> _SympyT: expr = expr.xreplace(min_max_replacements) if expr.has(TruncToInt): - trunc_replacements = {} + trunc_replacements: dict[sympy.Basic, sympy.Basic] = {} for atom in expr.atoms(TruncToInt): if isinstance(atom.args[0], IntTrueDiv): base, divisor = atom.args[0].args @@ -7008,7 +7027,7 @@ def simplify(self, expr: _SympyT, size_oblivious: bool = False) -> _SympyT: # for now just do a separate pass to catch common nested case if expr.has(FloorDiv): self._update_divisible() - div_replacements = {} + div_replacements: dict[sympy.Basic, sympy.Basic] = {} for atom in expr.atoms(FloorDiv): base, divisor = atom.args if isinstance(divisor, FloorDiv): @@ -7023,7 +7042,7 @@ def simplify(self, expr: _SympyT, size_oblivious: bool = False) -> _SympyT: expr = expr.xreplace(div_replacements) expr = safe_expand(expr) if expr.has(FloorDiv): - div_replacements = {} + div_replacements: dict[sympy.Basic, sympy.Basic] = {} pows = expr.atoms(sympy.Pow) rationals = expr.atoms(sympy.Rational).difference(expr.atoms(sympy.Integer)) for fd in expr.atoms(FloorDiv): @@ -7197,7 +7216,7 @@ def _make_data_dependent_error( def _update_var_to_range( self, symbol: sympy.Symbol, - vr: ValueRanges, + vr: ValueRanges[sympy.Expr], vr_sloc: ValueRangesSLoc | None = None, *, is_constraint: bool = False, @@ -7585,11 +7604,11 @@ def trivial_solve(lhs: sympy.Expr, rhs: sympy.Expr) -> bool: # See: Note - On 0/1 specialization def _default_value_range( self, do_not_specialize_zero_one: bool = False - ) -> ValueRanges: + ) -> ValueRanges[sympy.Expr]: lower = 0 if (do_not_specialize_zero_one or not self.specialize_zero_one) else 2 return ValueRanges(lower, int_oo) - def _default_unspecified_value_range(self) -> ValueRanges: + def _default_unspecified_value_range(self) -> ValueRanges[sympy.Expr]: return ValueRanges.unknown_int() @_lru_cache @@ -7840,8 +7859,8 @@ def evaluate_sym_node( self._expr_sym_node_id = id(sym_node) return self.evaluate_expr( sym_node.expr, - sym_node.hint, - sym_node.fx_node, + sym_node.hint, # pyrefly: ignore[bad-argument-type] + sym_node.fx_node, # pyrefly: ignore[bad-argument-type] size_oblivious, fallback_value=fallback_value, ) @@ -8688,6 +8707,8 @@ def _remove_effect_token_unbacked_bindings( # that we apply is unbacked renaming. def _get_placeholder_expr(sym_node: SymNode) -> sympy.Expr: shape_env = sym_node.shape_env + if shape_env is None: + raise AssertionError("shape_env is required for _get_placeholder_expr") result = sym_node._expr if result in shape_env.unbacked_renamings: return shape_env.unbacked_renamings[result] diff --git a/torch/fx/experimental/validator.py b/torch/fx/experimental/validator.py index 643e2d2537ef8..34046d1958580 100644 --- a/torch/fx/experimental/validator.py +++ b/torch/fx/experimental/validator.py @@ -327,7 +327,7 @@ def wrapper(*args): torch.sym_max: lift(ops.max), torch.sym_min: lift(ops.min), torch.sym_sum: lift(ops.sym_sum), - torch.sym_ite: lift(lambda b, t, f: t if b else f), + torch.sym_ite: lift(lambda b, t, f: z3.If(b, t, f)), torch._sym_sqrt: lift(ops.sqrt), # type: ignore[attr-defined] # Not lifted because we only use this function as a # marker for adding the expression as validator input. diff --git a/torch/fx/interpreter.py b/torch/fx/interpreter.py index a69e47e8d511c..7cd4ba1776567 100644 --- a/torch/fx/interpreter.py +++ b/torch/fx/interpreter.py @@ -1,8 +1,8 @@ -# mypy: allow-untyped-defs import inspect import logging +from collections.abc import Iterator from contextlib import contextmanager -from typing import Any, TYPE_CHECKING +from typing import Any import torch import torch.fx.traceback as fx_traceback @@ -19,15 +19,12 @@ from .proxy import Proxy -if TYPE_CHECKING: - from collections.abc import Iterator - log = logging.getLogger(__name__) __all__ = ["Interpreter", "Transformer"] -def _format_fx_node(n): +def _format_fx_node(n: Node) -> str: """ Format a torch.fx.Node into a human-readable string for debug logging. @@ -118,7 +115,7 @@ def __init__( module: torch.nn.Module, garbage_collect_values: bool = True, graph: Graph | None = None, - ): + ) -> None: self.module = module self.submodules = dict(self.module.named_modules()) if graph is not None: @@ -138,7 +135,7 @@ def __init__( node_to_last_use: dict[Node, Node] = {} self.user_to_last_uses: dict[Node, list[Node]] = {} - def register_last_uses(n: Node, user: Node): + def register_last_uses(n: Node, user: Node) -> None: if n not in node_to_last_use: node_to_last_use[n] = user self.user_to_last_uses.setdefault(user, []).append(n) @@ -150,7 +147,7 @@ def register_last_uses(n: Node, user: Node): @compatibility(is_backward_compatible=True) def run( self, - *args, + *args: Any, initial_env: dict[Node, Any] | None = None, enable_io_processing: bool = True, ) -> Any: @@ -240,7 +237,7 @@ def run( ) @compatibility(is_backward_compatible=True) - def boxed_run(self, args_list): + def boxed_run(self, args_list: list[Any]) -> Any: """ Run `module` via interpretation and return the result. This uses the "boxed" calling convention, where you pass a list of arguments, which will be cleared @@ -267,7 +264,7 @@ def boxed_run(self, args_list): return self.run(initial_env=env) @contextmanager - def _set_current_node(self, node): + def _set_current_node(self, node: Node) -> Iterator[None]: with fx_traceback.set_current_meta( node, f"Interpreter_{self.__class__.__name__}" ): @@ -452,7 +449,7 @@ def output( # Helper methods @compatibility(is_backward_compatible=True) - def fetch_attr(self, target: str): + def fetch_attr(self, target: str) -> Any: """ Fetch an attribute from the ``Module`` hierarchy of ``self.module``. @@ -473,7 +470,9 @@ def fetch_attr(self, target: str): return attr_itr @compatibility(is_backward_compatible=True) - def fetch_args_kwargs_from_env(self, n: Node) -> tuple[tuple, dict]: + def fetch_args_kwargs_from_env( + self, n: Node + ) -> tuple[tuple[Any, ...], dict[str, Any]]: """ Fetch the concrete values of ``args`` and ``kwargs`` of node ``n`` from the current execution environment. @@ -568,18 +567,18 @@ def fn(x): """ @compatibility(is_backward_compatible=True) - def __init__(self, module): + def __init__(self, module: GraphModule) -> None: super().__init__(module) self.new_graph = Graph() self.new_graph.set_codegen(module.graph._codegen) class TransformerTracer(Tracer): - def __init__(self, graph: Graph): + def __init__(self, graph: Graph) -> None: super().__init__() self.graph = graph self.tensor_attrs: dict[torch.Tensor, str] = {} # type: ignore[assignment] - def is_leaf_module(self, _, __) -> bool: + def is_leaf_module(self, _: torch.nn.Module, __: str) -> bool: return True self.tracer = TransformerTracer(self.new_graph) diff --git a/torch/fx/node.py b/torch/fx/node.py index 98ff8892243f2..ecb9fbd5ab27b 100644 --- a/torch/fx/node.py +++ b/torch/fx/node.py @@ -655,7 +655,7 @@ def format_node( return f"return {self.args[0]}" else: - def stringify_shape(shape: Iterable) -> str: + def stringify_shape(shape: Iterable[Any]) -> str: return f"[{', '.join([str(x) for x in shape])}]" meta_val = self.meta.get( diff --git a/torch/fx/operator_schemas.py b/torch/fx/operator_schemas.py index e2c76b6310260..532c8543ba888 100644 --- a/torch/fx/operator_schemas.py +++ b/torch/fx/operator_schemas.py @@ -1,4 +1,3 @@ -# mypy: allow-untyped-defs import enum import inspect import numbers @@ -6,11 +5,12 @@ import typing import warnings from collections.abc import Callable -from typing import Any, cast, NamedTuple, TYPE_CHECKING +from typing import Any, cast, Literal, NamedTuple, overload, TYPE_CHECKING import torch from torch._jit_internal import boolean_dispatched from torch._ops import OpOverload, OpOverloadPacket +from torch.utils._inspect import _fast_bind from ._compatibility import compatibility @@ -39,18 +39,18 @@ class ArgsKwargsPair(NamedTuple): kwargs: dict[str, Any] -_manual_overrides: dict[Callable, list[inspect.Signature]] = {} +_manual_overrides: dict[Callable[..., Any], list[inspect.Signature]] = {} -def _nonzero_schemas(): +def _nonzero_schemas() -> list[inspect.Signature]: signatures = [] - def nonzero(self): + def nonzero(self: torch.Tensor) -> None: pass signatures.append(inspect.signature(nonzero)) - def nonzero(self, *, as_tuple: bool): # type: ignore[no-redef] + def nonzero(self: torch.Tensor, *, as_tuple: bool) -> None: # type: ignore[no-redef] pass signatures.append(inspect.signature(nonzero)) @@ -62,7 +62,7 @@ def nonzero(self, *, as_tuple: bool): # type: ignore[no-redef] class _FakeGlobalNamespace: - def __getattr__(self, name): + def __getattr__(self, name: str) -> types.ModuleType: if name == "torch": return torch raise RuntimeError("Expected a torch namespace lookup") @@ -171,24 +171,26 @@ def _torchscript_schema_to_signature( @compatibility(is_backward_compatible=False) def check_for_mutable_operation( - target: Callable, args: tuple["Argument", ...], kwargs: dict[str, "Argument"] -): + target: Callable[..., Any], + args: tuple["Argument", ...], + kwargs: dict[str, "Argument"], +) -> None: signatures, schemas = get_signature_for_torch_op(target, return_schemas=True) if signatures and schemas: - matched_schemas = [] + matched_schemas: list[tuple[inspect.Signature, torch._C.FunctionSchema]] = [] # Iterate through all of the schema until we find one that matches # If one matches, populate `new_args_and_kwargs` with the new args/kwargs # values. If none matches, `new_args_and_kwargs` will be None for candidate_signature, schema in zip(signatures, schemas): try: - candidate_signature.bind(*args, **kwargs) + _fast_bind(candidate_signature, *args, **kwargs) matched_schemas.append((candidate_signature, schema)) except TypeError: continue - def throw_if_mutable(schema): + def throw_if_mutable(schema: torch._C.FunctionSchema) -> None: if schema.is_mutable: raise RuntimeError( f"Tried to trace mutable operation {schema}. FX only supports functional " @@ -209,8 +211,26 @@ def throw_if_mutable(schema): pass +@overload +def get_signature_for_torch_op( + op: Callable[..., Any], return_schemas: Literal[True] +) -> tuple[list[inspect.Signature] | None, list[torch._C.FunctionSchema] | None]: ... + + +@overload +def get_signature_for_torch_op( + op: Callable[..., Any], return_schemas: Literal[False] = ... +) -> list[inspect.Signature] | None: ... + + @compatibility(is_backward_compatible=False) -def get_signature_for_torch_op(op: Callable, return_schemas: bool = False): +def get_signature_for_torch_op( + op: Callable[..., Any], return_schemas: bool = False +) -> ( + list[inspect.Signature] + | tuple[list[inspect.Signature] | None, list[torch._C.FunctionSchema] | None] + | None +): """ Given an operator on the `torch` namespace, return a list of `inspect.Signature` objects corresponding to the overloads of that op.. May return `None` if a signature @@ -245,7 +265,7 @@ def get_signature_for_torch_op(op: Callable, return_schemas: bool = False): @compatibility(is_backward_compatible=False) -def create_type_hint(x): +def create_type_hint(x: object) -> object: """ Produces a type hint for the given argument. @@ -262,12 +282,12 @@ def create_type_hint(x): # todo(chilli): Figure out the right way for mypy to handle this if isinstance(x, list): - def ret_type(x): + def ret_type(x: Any) -> Any: return list[x] # type: ignore[valid-type] else: - def ret_type(x): + def ret_type(x: Any) -> Any: return tuple[x, ...] # type: ignore[valid-type] if len(x) == 0: @@ -290,7 +310,7 @@ def ret_type(x): @compatibility(is_backward_compatible=False) -def type_matches(signature_type: Any, argument_type: Any): +def type_matches(signature_type: Any, argument_type: Any) -> bool: sig_origin_type = getattr(signature_type, "__origin__", signature_type) if signature_type is argument_type: @@ -317,11 +337,11 @@ def type_matches(signature_type: Any, argument_type: Any): if getattr(argument_type, "__origin__", None) is list: return issubclass(argument_type.__args__[0], sig_el_type) - def is_homogeneous_tuple(t): - if getattr(t, "__origin__", None) is not tuple: + def is_homogeneous_tuple(t: object) -> bool: + if typing.get_origin(t) is not tuple: return False - contained = t.__args__ - if t.__args__ == ((),): # Tuple[()].__args__ == ((),) for some reason + contained = typing.get_args(t) + if contained == ((),): # Tuple[()].__args__ == ((),) for some reason return True return all((c is Ellipsis) or issubclass(c, sig_el_type) for c in contained) @@ -342,7 +362,7 @@ def is_homogeneous_tuple(t): @compatibility(is_backward_compatible=False) def _normalize_function_or_error( - target: Callable, + target: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any] | None = None, arg_types: tuple[Any] | None = None, @@ -366,7 +386,7 @@ def _normalize_function_or_error( @compatibility(is_backward_compatible=False) def normalize_function( - target: Callable, + target: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any] | None = None, arg_types: tuple[Any] | None = None, @@ -442,14 +462,14 @@ def normalize_function( if not callable(target): raise AssertionError(f"target must be callable, got {type(target)}") torch_op_schemas = get_signature_for_torch_op(target) - matched_schemas = [] + matched_schemas: list[inspect.Signature] = [] if torch_op_schemas: # Iterate through all of the schema until we find one that matches # If one matches, populate `new_args_and_kwargs` with the new args/kwargs # values. If none matches, `new_args_and_kwargs` will be None for candidate_signature in torch_op_schemas: try: - candidate_signature.bind(*args, **kwargs) + _fast_bind(candidate_signature, *args, **kwargs) matched_schemas.append(candidate_signature) except TypeError: continue @@ -469,8 +489,8 @@ def normalize_function( for candidate_signature in torch_op_schemas: sig_matches = True try: - bound_types = candidate_signature.bind( - *arg_types, **kwarg_types + bound_types = _fast_bind( + candidate_signature, *arg_types, **kwarg_types ) for arg_name, arg_type in bound_types.arguments.items(): param = candidate_signature.parameters[arg_name] @@ -589,7 +609,7 @@ def _args_kwargs_to_normalized_args_kwargs( if list(sig.parameters.keys()) != ["input", "from", "to", "generator"]: return None - bound_args = sig.bind(*args, **kwargs) + bound_args = _fast_bind(sig, *args, **kwargs) bound_args.apply_defaults() new_kwargs: dict[str, Any] = {} diff --git a/torch/fx/passes/_tensorify_python_scalars.py b/torch/fx/passes/_tensorify_python_scalars.py index 253521435f7f4..c70da6a67c8c0 100644 --- a/torch/fx/passes/_tensorify_python_scalars.py +++ b/torch/fx/passes/_tensorify_python_scalars.py @@ -347,7 +347,7 @@ def _sympy_interp(expr: sympy.Expr) -> MetaProxy: ): failed_tensorify_ops.update(str(node.target)) - log.info("Failed to tensorify %s", str(node.target)) + log.info("Failed to tensorify %s", node.target) # Now do one more pass that specializes all symfloats we didn't manage # to tensorify away. diff --git a/torch/fx/passes/backends/cudagraphs.py b/torch/fx/passes/backends/cudagraphs.py index 97496fbc9b2a2..ec9cb02c43ff9 100644 --- a/torch/fx/passes/backends/cudagraphs.py +++ b/torch/fx/passes/backends/cudagraphs.py @@ -1,7 +1,11 @@ -# mypy: allow-untyped-defs import operator +from collections.abc import Mapping, Sequence +from typing import Any import torch + + +__all__ = ["CudaGraphsSupport", "partition_cudagraphs"] from torch.fx.passes.fake_tensor_prop import FakeTensorProp from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner from torch.fx.passes.operator_support import OperatorSupport @@ -11,7 +15,9 @@ class CudaGraphsSupport(OperatorSupport): # TODO: why is submodules passed here - def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: + def is_node_supported( + self, submodules: Mapping[str, torch.nn.Module], node: torch.fx.Node + ) -> bool: if node.op not in CALLABLE_NODE_OPS: return False @@ -23,10 +29,10 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: found_not_cuda = False - def meta_fk(meta): + def meta_fk(meta: dict[str, Any]) -> torch.Tensor: return meta["val"] if "val" in meta else meta["fake_result"] - def find_not_cuda(t): + def find_not_cuda(t: object) -> None: nonlocal found_not_cuda if isinstance(t, torch.Tensor) and t.device.type != "cuda": found_not_cuda = True @@ -42,7 +48,9 @@ def find_not_cuda(t): return not found_not_cuda -def partition_cudagraphs(gm, inputs): +def partition_cudagraphs( + gm: torch.fx.GraphModule, inputs: Sequence[object] +) -> torch.fx.GraphModule: """ Partition an FX graph into sub-GraphModules that can be validly run under CUDA graphs. For a subgraph to be runnable under CUDA, all of the operations diff --git a/torch/fx/passes/dialect/common/cse_pass.py b/torch/fx/passes/dialect/common/cse_pass.py index e5889375bb07a..c5c631e53c0b5 100644 --- a/torch/fx/passes/dialect/common/cse_pass.py +++ b/torch/fx/passes/dialect/common/cse_pass.py @@ -1,8 +1,8 @@ -# mypy: allow-untyped-defs from typing import Any import torch from torch.fx import Graph, GraphModule, Node +from torch.fx.node import Target from torch.fx.passes.infra.pass_base import PassBase, PassResult from torch.utils._pytree import tree_flatten @@ -43,13 +43,15 @@ @torch.fx._compatibility.compatibility(is_backward_compatible=False) -def get_CSE_banned_ops(): +def get_CSE_banned_ops() -> set[torch._ops.OpOverloadPacket]: return rand_ops.union(inplace_ops) @torch.fx._compatibility.compatibility(is_backward_compatible=False) class CSEPass(PassBase): - def __init__(self, banned_ops=None): + def __init__( + self, banned_ops: set[torch._ops.OpOverloadPacket] | None = None + ) -> None: """ This version of CSE Pass aims to be dialect agnostic, and it's implemented purely based on the connectivity between fx.Node. @@ -83,7 +85,7 @@ def f(a): print(result.graph_module) """ - def get_aten_target(node): + def get_aten_target(node: Node) -> Target: if hasattr(node.target, "overloadpacket"): return node.target.overloadpacket return node.target @@ -94,10 +96,10 @@ def get_aten_target(node): Node, Node ] = {} # map from node in the old graph to node in the new graph hash_env: dict[ - tuple[torch._ops.OpOverload, int], Node + tuple[Target, int], Node ] = {} # map from hash to a node in the new graph token_map: dict[ - tuple[torch._ops.OpOverload, int], dict[str, Any] + tuple[Target, int], dict[str, Any] ] = {} # map from hash to token for n in graph_module.graph.nodes: # The placeholder, output, and get_attr nodes are copied to the new graph without change @@ -113,7 +115,7 @@ def get_aten_target(node): else: # n.op == 'call_function', should never see n.op == 'call_module' or 'call_method' # substitute args and kwargs members to their mapping in env if exists # specs can be used to reconstruct nested list/dictionaries - def substitute(arg_list): + def substitute(arg_list: Any) -> tuple[tuple[Any, ...], Any]: arg_list, spec = tree_flatten(arg_list) for i in range(len(arg_list)): v = arg_list[i] diff --git a/torch/fx/passes/fake_tensor_prop.py b/torch/fx/passes/fake_tensor_prop.py index 45507632228ef..f33e3eeb8ea9a 100644 --- a/torch/fx/passes/fake_tensor_prop.py +++ b/torch/fx/passes/fake_tensor_prop.py @@ -1,4 +1,4 @@ -# mypy: allow-untyped-defs +from typing import Any import torch.fx from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode @@ -29,7 +29,7 @@ class FakeTensorProp(torch.fx.Interpreter): def __init__( self, module: torch.fx.GraphModule, mode: FakeTensorMode | None = None - ): + ) -> None: super().__init__(module) if mode is None: mode = FakeTensorMode() @@ -38,7 +38,7 @@ def __init__( mode.reset_nt_tensor_id_counter() self.seen_subgraphs: OrderedSet[str] = OrderedSet() - def run_node(self, n: Node): + def run_node(self, n: Node) -> Any: from torch.fx.experimental.symbolic_shapes import ( compute_unbacked_bindings, rebind_unbacked, @@ -79,7 +79,7 @@ def run_node(self, n: Node): result = super().run_node(n) rebind_unbacked(self._mode.shape_env, n, result) - def extract_val(obj): + def extract_val(obj: Any) -> Any: if isinstance(obj, FakeTensor): return snapshot_fake(obj) elif isinstance(obj, torch.Tensor): @@ -101,13 +101,13 @@ def extract_val(obj): return result - def propagate(self, *args): + def propagate(self, *args: object) -> object: fake_args = [ self._mode.from_tensor(a) if isinstance(a, torch.Tensor) else a for a in args ] return self.propagate_dont_convert_inputs(*fake_args) - def propagate_dont_convert_inputs(self, *args): + def propagate_dont_convert_inputs(self, *args: object) -> object: with self._mode: return super().run(*args) diff --git a/torch/fx/passes/graph_drawer.py b/torch/fx/passes/graph_drawer.py index f94d4cd174f53..8687be9126b3d 100644 --- a/torch/fx/passes/graph_drawer.py +++ b/torch/fx/passes/graph_drawer.py @@ -1,5 +1,3 @@ -# mypy: allow-untyped-defs - import hashlib from itertools import chain from types import ModuleType @@ -85,7 +83,7 @@ def __init__( parse_stack_trace: bool = False, dot_graph_shape: str | None = None, normalize_args: bool = False, - ): + ) -> None: self._name = name self.dot_graph_shape = ( dot_graph_shape if dot_graph_shape is not None else "record" @@ -122,7 +120,7 @@ def __init__( parse_stack_trace, ) - def get_dot_graph(self, submod_name=None) -> pydot.Dot: + def get_dot_graph(self, submod_name: str | None = None) -> pydot.Dot: """ Visualize a torch.fx.Graph with graphviz Example: @@ -154,7 +152,7 @@ def get_dot_graph(self, submod_name=None) -> pydot.Dot: def get_main_dot_graph(self) -> pydot.Dot: return self._dot_graphs[self._name] - def get_submod_dot_graph(self, submod_name) -> pydot.Dot: + def get_submod_dot_graph(self, submod_name: str) -> pydot.Dot: return self._dot_graphs[f"{self._name}_{submod_name}"] def get_all_dot_graphs(self) -> dict[str, pydot.Dot]: @@ -198,7 +196,7 @@ def _get_leaf_node( py_obj = getattr(py_obj, atom) return py_obj - def _typename(self, target: Any) -> str: + def _typename(self, target: torch.fx.node.Target | torch.nn.Module) -> str: if isinstance(target, torch.nn.Module): ret = torch.typename(target) elif isinstance(target, str): @@ -218,7 +216,7 @@ def _shorten_file_name( self, full_file_name: str, truncate_to_last_n: int = 2, - ): + ) -> str: splits = full_file_name.split("/") if len(splits) >= truncate_to_last_n: return "/".join(splits[-truncate_to_last_n:]) @@ -231,7 +229,7 @@ def _get_node_label( skip_node_names_in_args: bool, parse_stack_trace: bool, ) -> str: - def _get_str_for_args_kwargs(arg): + def _get_str_for_args_kwargs(arg: tuple[Any, ...] | dict[str, Any]) -> str: if isinstance(arg, tuple): prefix, suffix = r"|args=(\l", r",\n)\l" arg_strs_list = [_format_arg(a, max_list_len=8) for a in arg] @@ -311,7 +309,7 @@ def _get_str_for_args_kwargs(arg): return label + "}" - def _tensor_meta_to_label(self, tm) -> str: + def _tensor_meta_to_label(self, tm: object) -> str: if tm is None: return "" elif isinstance(tm, TensorMetadata): @@ -414,7 +412,7 @@ def _to_dot( # "TB" means top-to-bottom rank direction in layout dot_graph = pydot.Dot(name, rankdir="TB") - buf_name_to_subgraph = {} + buf_name_to_subgraph: dict[str, pydot.Cluster] = {} for node in graph_module.graph.nodes: if ignore_getattr and node.op == "get_attr": @@ -443,7 +441,7 @@ def _to_dot( # pyrefly: ignore [missing-attribute] current_graph.add_node(dot_node) - def get_module_params_or_buffers(): + def get_module_params_or_buffers() -> None: for pname, ptensor in chain( leaf_module.named_parameters(), # pyrefly: ignore [bad-argument-type] diff --git a/torch/fx/passes/graph_manipulation.py b/torch/fx/passes/graph_manipulation.py index 64e1481915e00..222245718b9c4 100644 --- a/torch/fx/passes/graph_manipulation.py +++ b/torch/fx/passes/graph_manipulation.py @@ -1,12 +1,11 @@ -# mypy: allow-untyped-defs -from typing import Any, NamedTuple +from typing import NamedTuple import torch from torch.fx._compatibility import compatibility from torch.fx.graph import Graph from torch.fx.graph_module import GraphModule from torch.fx.node import map_arg, Node, Target -from torch.fx.passes.shape_prop import ShapeProp +from torch.fx.passes.shape_prop import ShapeProp, TensorMetadata __all__ = [ @@ -25,9 +24,11 @@ def replace_target_nodes_with( old_target: Target, new_op: str, new_target: Target, -): - """Modifies all nodes in fx_module.graph.nodes which match the specified op code and target, - and updates them to match the new op code and target""" +) -> None: + """ + Modifies all nodes in fx_module.graph.nodes which match the specified op code + and target, and updates them to match the new op code and target. + """ new_graph = Graph() val_map: dict[Node, Node] = {} for node in fx_module.graph.nodes: @@ -71,7 +72,7 @@ def get_size_of_all_nodes( @compatibility(is_backward_compatible=False) -def get_tensor_meta(node: Node) -> Any: +def get_tensor_meta(node: Node) -> TensorMetadata: tensor_meta = node.meta.get("tensor_meta") if not tensor_meta: diff --git a/torch/fx/passes/graph_transform_observer.py b/torch/fx/passes/graph_transform_observer.py index f91877b00a057..e5e5e5c3aaf5b 100644 --- a/torch/fx/passes/graph_transform_observer.py +++ b/torch/fx/passes/graph_transform_observer.py @@ -1,6 +1,6 @@ -# mypy: allow-untyped-defs import os from collections.abc import Callable +from types import TracebackType from typing import TypeVar from torch.fx import Graph, Node @@ -28,7 +28,7 @@ def __init__( passname: str, subsystem: str | None = None, log_url: str | None = None, - ): + ) -> None: """ log_url is inferred to be torch._inductor.config.trace.log_url_for_graph_xform unless otherwise specified """ @@ -73,7 +73,7 @@ def __init__( ).get_dot_graph() @classmethod - def get_current_pass_count(cls): + def get_current_pass_count(cls) -> int: return cls.__pass_count def apply_gm_pass(self, pass_fn: Callable[[GraphModule], T]) -> T | None: @@ -102,7 +102,7 @@ def apply_graph_pass(self, pass_fn: Callable[[Graph], T]) -> T | None: ): return pass_fn(self.gm.graph) - def _check_disable_pass(self): + def _check_disable_pass(self) -> bool: from torch._inductor import config as inductor_config if self.passname.upper() in inductor_config.disabled_passes.upper(): @@ -118,7 +118,7 @@ def _check_disable_pass(self): "inductor", self.subsystem, debug_info ) - def __enter__(self): + def __enter__(self) -> "GraphTransformObserver": if not self.active: return self self.gm._register_create_node_hook(self._node_creation_hook) @@ -136,7 +136,12 @@ def __enter__(self): return self - def __exit__(self, type, value, tb): + def __exit__( + self, + type: type[BaseException] | None, + value: BaseException | None, + tb: TracebackType | None, + ) -> None: if not self.active: return for gm in self.copied_gms + [self.gm]: @@ -181,10 +186,10 @@ def __exit__(self, type, value, tb): ) ) - def get_node_creation_hook(self): + def get_node_creation_hook(self) -> Callable[[Node], None]: # We have to return a function instead of using a class method directly # to avoid max recursion issue when deepcopy a graph module within the context manager. - def on_node_creation(node): + def on_node_creation(node: Node) -> None: self.created_nodes.add(node.name) self.name_to_node[node.name] = node source = NodeSource(None, self.passname, NodeSourceAction.CREATE) @@ -195,15 +200,15 @@ def on_node_creation(node): return on_node_creation - def get_node_erase_hook(self): - def on_node_erase(node): + def get_node_erase_hook(self) -> Callable[[Node], None]: + def on_node_erase(node: Node) -> None: self.erased_nodes.add(node.name) self.name_to_node.pop(node.name, None) return on_node_erase - def get_node_replace_hook(self): - def on_node_replace(old: Node, new: str, user: Node): + def get_node_replace_hook(self) -> Callable[[Node, str, Node], None]: + def on_node_replace(old: Node, new: str, user: Node) -> None: # Update node meta when replacing old node with new node new_node = self.name_to_node.get(new, None) @@ -223,7 +228,7 @@ def on_node_replace(old: Node, new: str, user: Node): if new_node.name in self.created_nodes: action.append(NodeSourceAction.CREATE) - def created_this_pass(source): + def created_this_pass(source: NodeSource) -> bool: return source.pass_name == self.passname and source.action == [ NodeSourceAction.CREATE ] @@ -241,8 +246,8 @@ def created_this_pass(source): return on_node_replace - def get_deepcopy_hook(self): - def on_deepcopy(gm): + def get_deepcopy_hook(self) -> Callable[[GraphModule], None]: + def on_deepcopy(gm: GraphModule) -> None: self.copied_gms.append(gm) return on_deepcopy diff --git a/torch/fx/passes/infra/partitioner.py b/torch/fx/passes/infra/partitioner.py index 550c6d57be865..be8b4575a4c4e 100644 --- a/torch/fx/passes/infra/partitioner.py +++ b/torch/fx/passes/infra/partitioner.py @@ -1,4 +1,3 @@ -# mypy: allow-untyped-defs import collections import itertools import logging @@ -21,7 +20,7 @@ def __init__( id: int | None = None, nodes: Iterable[Node] | None = None, node_orders: Iterable[int] | None = None, - ): + ) -> None: self.id = id self.nodes: dict[Node, int | None] = {} if nodes is not None: @@ -39,18 +38,18 @@ def __init__( def __repr__(self) -> str: return str(self.nodes) - def add_node(self, node: Node, node_order: int | None = None): + def add_node(self, node: Node, node_order: int | None = None) -> None: self.nodes.update({node: node_order}) - def remove_node(self, node: Node): + def remove_node(self, node: Node) -> None: del self.nodes[node] - def size(self): + def size(self) -> int: return len(self.nodes) class _DependencyViewer: - def __init__(self, graph_module: GraphModule): + def __init__(self, graph_module: GraphModule) -> None: self.downstreams = collections.defaultdict(set) for node in reversed(graph_module.graph.nodes): @@ -92,7 +91,7 @@ def propose_partitions(self) -> list[Partition]: # partition_map is a mapping from partition id to a set of partition id's. # The value set contains all the partition ids that can be reached by doing a # DFS starting from the partition id in the key. - partition_map: dict[int, set] = collections.defaultdict(set) + partition_map: dict[int, set[int]] = collections.defaultdict(set) # assumptions: nodes in candidate list is sorted in topological order assignment: dict[Node, int] = {} # mapping from node to partition_id @@ -106,19 +105,19 @@ def propose_partitions(self) -> list[Partition]: int, int ] = {} # mapping from partition_id to minimum topo order of nodes in partition partition_users: dict[ - int, set + int, set[Node] ] = {} # mapping from partition_id to partition users new_partition_id = itertools.count() # try to merge partition other_id into partition self_id # merge only happens if the end graph doesn't contain cyclic dependency # returns `True` when merge happens, `False` otherwise. - def maybe_merge_partition(self_id: int, other_id: int): + def maybe_merge_partition(self_id: int, other_id: int) -> tuple[int, bool]: # merged_nodes is the union of nodes in two partition to-be-merged self_nodes = partitions_by_id[self_id].nodes other_nodes = partitions_by_id[other_id].nodes - def dfs_iter_find_cycle(all_user_nodes: set[Node]): + def dfs_iter_find_cycle(all_user_nodes: set[Node]) -> bool: for user_node in all_user_nodes: visited_partition_ids = set() @@ -185,8 +184,10 @@ def dfs_iter_find_cycle(all_user_nodes: set[Node]): return merge_id, True - def merge_single_node(node: Node, node_order: int | None, id: int | None): - def _update_partition_map(node: Node, id: int): + def merge_single_node( + node: Node, node_order: int | None, id: int | None + ) -> None: + def _update_partition_map(node: Node, id: int) -> None: # Iterate through all the users of this node and update the partition map to indicate # that there is a path from the partition id of this node to the target partition id. for user_node in node.users: @@ -331,10 +332,10 @@ def fuse_partitions( ) # remove non-compute-ops that sits at the boundary of a partition. - def remove_bookend_non_compute_ops(self, partitions: list[Partition]): + def remove_bookend_non_compute_ops(self, partitions: list[Partition]) -> None: non_compute_ops = set(self.non_compute_ops) - def is_non_compute_node(node: Node): + def is_non_compute_node(node: Node) -> bool: return ( node.op == "call_function" and _get_qualified_name(node.target) in non_compute_ops # type: ignore[arg-type] @@ -346,7 +347,7 @@ def is_non_compute_node(node: Node): def is_transparent_input_node( node: Node, partition: set[Node], removed_nodes: set[Node] - ): + ) -> bool: if ( node.op == "placeholder" or (node not in partition) @@ -367,7 +368,7 @@ def is_transparent_input_node( def is_transparent_output_node( node: Node, partition: set[Node], removed_nodes: set[Node] - ): + ) -> bool: if ( node.op == "placeholder" or (node not in partition) diff --git a/torch/fx/passes/infra/pass_base.py b/torch/fx/passes/infra/pass_base.py index 109f62f1488bb..24571d82cd4a1 100644 --- a/torch/fx/passes/infra/pass_base.py +++ b/torch/fx/passes/infra/pass_base.py @@ -1,7 +1,7 @@ -# mypy: allow-untyped-defs import abc from collections import namedtuple +import torch.nn as nn from torch.fx._compatibility import compatibility from torch.fx.graph_module import GraphModule @@ -20,7 +20,7 @@ class PassResult(namedtuple("PassResult", ["graph_module", "modified"])): __slots__ = () - def __new__(cls, graph_module, modified): + def __new__(cls, graph_module: nn.Module, modified: bool) -> "PassResult": return super().__new__(cls, graph_module, modified) diff --git a/torch/fx/passes/infra/pass_manager.py b/torch/fx/passes/infra/pass_manager.py index 87fb6e70037f9..a474d692627c0 100644 --- a/torch/fx/passes/infra/pass_manager.py +++ b/torch/fx/passes/infra/pass_manager.py @@ -1,9 +1,9 @@ -# mypy: allow-untyped-defs import inspect import logging from collections.abc import Callable from functools import wraps from queue import Queue +from typing import Any import torch.nn as nn from torch.fx._compatibility import compatibility @@ -18,7 +18,7 @@ @compatibility(is_backward_compatible=False) -def pass_result_wrapper(fn: Callable) -> Callable: +def pass_result_wrapper(fn: Callable[..., Any]) -> Callable[..., PassResult | None]: """ Wrapper for passes which currently do not return a PassResult. This wrapper makes them return a PassResult containing the modified object @@ -35,7 +35,7 @@ def pass_result_wrapper(fn: Callable) -> Callable: return None @wraps(fn) - def wrapped_fn(gm): + def wrapped_fn(gm: nn.Module) -> PassResult | None: res = fn(gm) if res is None: return PassResult(gm, True) @@ -51,7 +51,8 @@ def wrapped_fn(gm): def _validate_pass_schedule_constraint( - constraint: Callable[[Callable, Callable], bool], passes: list[Callable] + constraint: Callable[[Callable[..., Any], Callable[..., Any]], bool], + passes: list[Callable[..., Any]], ) -> None: for i, a in enumerate(passes): for j, b in enumerate(passes[i + 1 :]): @@ -65,8 +66,8 @@ def _validate_pass_schedule_constraint( def _topological_sort_passes( - passes: list[Callable], constraints: list[Callable] -) -> list[Callable]: + passes: list[Callable[..., Any]], constraints: list[Callable[..., Any]] +) -> list[Callable[..., Any]]: """ Args passes: Passes that we are ordering @@ -80,9 +81,9 @@ def _topological_sort_passes( return passes # Construct a graph mapping nodes to a list of their users - graph: dict[Callable, list[Callable]] = {p: [] for p in passes} - indegree_map: dict[Callable, int] = dict.fromkeys(passes, 0) - candidates: Queue = Queue() + graph: dict[Callable[..., Any], list[Callable[..., Any]]] = {p: [] for p in passes} + indegree_map: dict[Callable[..., Any], int] = dict.fromkeys(passes, 0) + candidates: Queue[Callable[..., Any]] = Queue() for a in passes: for b in passes: if a == b: @@ -96,8 +97,8 @@ def _topological_sort_passes( if indegree_map[a] == 0: candidates.put(a) - visited: dict[Callable, bool] = dict.fromkeys(passes, False) - sorted_passes: list[Callable] = [] + visited: dict[Callable[..., Any], bool] = dict.fromkeys(passes, False) + sorted_passes: list[Callable[..., Any]] = [] while not candidates.empty(): p = candidates.get() @@ -122,27 +123,28 @@ def _topological_sort_passes( @compatibility(is_backward_compatible=False) -def this_before_that_pass_constraint(this: Callable, that: Callable) -> Callable: +def this_before_that_pass_constraint( + this: Callable[..., Any], that: Callable[..., Any] +) -> Callable[[Callable[..., Any], Callable[..., Any]], bool]: """ - Defines a partial order ('depends on' function) where `this` must occur - before `that`. + Defines a partial order ('depends on' function) where ``this`` must occur + before ``that``. - For example, the following pass list and constraint list would be invalid. - ``` - passes = [pass_b, pass_a] + For example, the following pass list and constraint list would be invalid:: - constraints = [this_before_that_pass_constraint(pass_a, pass_b)] - ``` + passes = [pass_b, pass_a] + + constraints = [this_before_that_pass_constraint(pass_a, pass_b)] Args: this (Callable): pass which should occur first that (Callable): pass which should occur later Returns: - depends_on (Callable[[Object, Object], bool] + depends_on (Callable[[Object, Object], bool]) """ - def depends_on(a: Callable, b: Callable): + def depends_on(a: Callable[..., Any], b: Callable[..., Any]) -> bool: return a != that or b != this return depends_on @@ -170,19 +172,20 @@ class PassManager: checks """ - passes: list[Callable[[nn.Module], PassResult]] - constraints: list[Callable[[Callable, Callable], bool]] + passes: list[Callable[..., PassResult | None]] + constraints: list[Callable[[Callable[..., Any], Callable[..., Any]], bool]] _validated: bool = False steps: int = 1 def __init__( self, - passes=None, - constraints=None, - steps=None, + passes: list[Callable[..., PassResult | None]] | None = None, + constraints: list[Callable[[Callable[..., Any], Callable[..., Any]], bool]] + | None = None, + steps: int | None = None, run_checks_after_each_pass: bool = False, suppress_check_failures: bool = False, - ): + ) -> None: self.passes = passes or [] self.constraints = constraints or [] if steps: @@ -191,21 +194,23 @@ def __init__( self.run_checks_after_each_pass = run_checks_after_each_pass self.suppress_check_failures = suppress_check_failures - def add_pass(self, _pass: Callable): + def add_pass(self, _pass: Callable[..., PassResult | None]) -> None: """ Adds a pass into the current list of passes. """ self.passes.append(_pass) self._validated = False - def add_constraint(self, constraint: Callable): + def add_constraint( + self, constraint: Callable[[Callable[..., Any], Callable[..., Any]], bool] + ) -> None: """ Adds a constraint into the current list of constraints. """ self.constraints.append(constraint) self._validated = False - def validate_constraints(self): + def validate_constraints(self) -> None: """ Validates that current pass schedule defined by `self.passes` is valid according to all constraints in `self.constraints` @@ -216,7 +221,7 @@ def validate_constraints(self): _validate_pass_schedule_constraint(constraint, self.passes) self._validated = True - def solve_constraints(self): + def solve_constraints(self) -> None: """ Finds a valid traversal order based on the given constraints and orders the passes based on this order. @@ -228,7 +233,7 @@ def solve_constraints(self): self.passes = _topological_sort_passes(self.passes, self.constraints) self._validated = True - def add_checks(self, check: Callable) -> None: + def add_checks(self, check: Callable[[nn.Module], None]) -> None: """ Adds a function which takes runs various checks on a given graph module. This function is run before and after each pass if the @@ -283,7 +288,9 @@ def __call__(self, module: nn.Module) -> PassResult: f"The result of the pass {fn_name} should be type PassResult." + "Please wrap it with pass_result_wrapper()" ) + # pyrefly: ignore[missing-attribute] module = res.graph_module + # pyrefly: ignore[missing-attribute] modified = modified or res.modified if isinstance(module, GraphModule): diff --git a/torch/fx/passes/net_min_base.py b/torch/fx/passes/net_min_base.py index 34636c6f3664a..7335497a670ef 100644 --- a/torch/fx/passes/net_min_base.py +++ b/torch/fx/passes/net_min_base.py @@ -1,4 +1,3 @@ -# mypy: allow-untyped-defs import logging from collections.abc import Callable from dataclasses import dataclass @@ -78,7 +77,7 @@ class _MinimizerSettingBase: return_intermediate: bool = False all_outputs: bool = False - def __str__(self): + def __str__(self) -> str: settings_str = "FX Minimizer Settings:\n" for k, v in vars(self).items(): @@ -113,7 +112,7 @@ def __init__( module_exporter: Callable[[Tensors, torch.fx.GraphModule, str], None] | None = None, exclusion_fn: Callable[[NodeList, int, int], None] | None = None, - ): + ) -> None: if not isinstance(module, torch.fx.GraphModule): raise AssertionError(f"Expected GraphModule, got {type(module)}") @@ -190,7 +189,7 @@ def _store_outputs( a_result: TensorOrTensors, b_result: TensorOrTensors, submodule: torch.fx.GraphModule, - ): + ) -> None: """ Store the outputs of self.run_a() and self.run_b() into self.a_outputs and self.b_outputs, so that we can use them when execute preceding nodes that @@ -251,7 +250,7 @@ def _get_submod_inputs( if self.settings.accumulate_error: print(f"Can't find previous stored outputs named {placeholders}!") - def get_inputs(self: torch.nn.Module, inputs: Any): + def get_inputs(self: torch.nn.Module, inputs: tuple[Any, ...]) -> None: nonlocal a_input a_input = inputs @@ -267,7 +266,7 @@ def get_inputs(self: torch.nn.Module, inputs: Any): return a_input, b_input - def _tag_nodes(self, selected_nodes: NodeSet): + def _tag_nodes(self, selected_nodes: NodeSet) -> None: """ Tag selected nodes with tag "minimize". Nodes with the same tags will be split to the same submodule afterwards. @@ -336,7 +335,7 @@ def _run_and_compare( submod_name: str, output_names: Names, report_idx: int = -1, - ): + ) -> None: """ Run the submodule in `split_module` that has name `submod_name` using `self.run_a` and `self.run_b` and compare their results. @@ -837,7 +836,7 @@ def _skip_traverse_impl( self.print_report(report) return set() - def _skip_traverse(self, all_nodes: NodeList, skip_nodes: list) -> NodeSet: + def _skip_traverse(self, all_nodes: NodeList, skip_nodes: list[str]) -> NodeSet: """ Skip certain nodes in graph based on settings """ @@ -880,7 +879,7 @@ def _collect_nodes(self, start: str | None, end: str | None) -> NodeList: return nodes - def run_nodes(self, start: str | None = None, end: str | None = None): + def run_nodes(self, start: str | None = None, end: str | None = None) -> None: """ Run part of the model from `start` node to `end` node. If `start` is None then we start from the beginning of the model. If `end` is None then we @@ -901,7 +900,7 @@ def run_nodes(self, start: str | None = None, end: str | None = None): if node in self.fusions: cur_nodes.update(self.fusions[node]) - output_names = [] + output_names: list[str] = [] if self.settings.return_intermediate: output_names = [node.name for node in nodes] @@ -914,14 +913,14 @@ def run_nodes(self, start: str | None = None, end: str | None = None): ) as e: print(e) - def print_report(self, report: list[str]): + def print_report(self, report: list[str]) -> None: for i in range(len(report)): if i > 0: print(" . " + report[i]) else: print(report[i]) - def print_reports(self): + def print_reports(self) -> None: for report in self.reports: self.print_report(report) @@ -929,7 +928,7 @@ def minimize( self, start: str | None = None, end: str | None = None, - skip_nodes: list | None = None, + skip_nodes: list[str] | None = None, find_last_node: bool | None = None, ) -> NodeSet: """ diff --git a/torch/fx/passes/operator_support.py b/torch/fx/passes/operator_support.py index 813e22090ba7a..64d40737dc275 100644 --- a/torch/fx/passes/operator_support.py +++ b/torch/fx/passes/operator_support.py @@ -1,4 +1,3 @@ -# mypy: allow-untyped-defs import abc import typing as t @@ -67,7 +66,7 @@ class OperatorSupport(OperatorSupportBase): _support_dict: SupportDict - def __init__(self, support_dict: SupportDict | None = None): + def __init__(self, support_dict: SupportDict | None = None) -> None: self._support_dict = support_dict or {} def is_node_supported( @@ -163,7 +162,7 @@ def chain(*op_support: OperatorSupportBase) -> OperatorSupportBase: any of it reports False. """ - def _chain(submods, node) -> bool: + def _chain(submods: t.Mapping[str, torch.nn.Module], node: torch.fx.Node) -> bool: return all(x.is_node_supported(submods, node) for x in op_support) return create_op_support(_chain) @@ -176,7 +175,9 @@ def any_chain(*op_support: OperatorSupportBase) -> OperatorSupportBase: any of it reports True. """ - def _any_chain(submods, node) -> bool: + def _any_chain( + submods: t.Mapping[str, torch.nn.Module], node: torch.fx.Node + ) -> bool: return any(x.is_node_supported(submods, node) for x in op_support) return create_op_support(_any_chain) @@ -219,7 +220,7 @@ def _decline_if_node_in_names( return create_op_support(_decline_if_node_in_names) -def _get_arg_dtype(arg: torch.fx.Node) -> t.Any: +def _get_arg_dtype(arg: torch.fx.Node) -> torch.dtype: if not isinstance(arg, torch.fx.Node): raise AssertionError(f"Expected torch.fx.Node, got {type(arg)}") tensor_meta = arg.meta.get("tensor_meta") # type: ignore[union-attr] diff --git a/torch/fx/passes/pass_manager.py b/torch/fx/passes/pass_manager.py index 918f628d73eff..b413a83744beb 100644 --- a/torch/fx/passes/pass_manager.py +++ b/torch/fx/passes/pass_manager.py @@ -1,8 +1,13 @@ -# mypy: allow-untyped-defs import logging from collections.abc import Callable from functools import wraps from inspect import unwrap +from typing import Any, Concatenate, ParamSpec, TypeVar + + +_P = ParamSpec("_P") +_R = TypeVar("_R") +_T = TypeVar("_T") logger = logging.getLogger(__name__) @@ -19,7 +24,9 @@ # for callables which modify object inplace and return something other than # the object on which they act -def inplace_wrapper(fn: Callable) -> Callable: +def inplace_wrapper( + fn: Callable[Concatenate[_T, _P], Any], +) -> Callable[Concatenate[_T, _P], _T]: """ Convenience wrapper for passes which modify an object inplace. This wrapper makes them return the modified object instead. @@ -32,33 +39,32 @@ def inplace_wrapper(fn: Callable) -> Callable: """ @wraps(fn) - def wrapped_fn(gm): - fn(gm) + def wrapped_fn(gm: _T, *args: _P.args, **kwargs: _P.kwargs) -> _T: + fn(gm, *args, **kwargs) return gm return wrapped_fn -def log_hook(fn: Callable, level=logging.INFO) -> Callable: +def log_hook(fn: Callable[_P, _R], level: int = logging.INFO) -> Callable[_P, _R]: """ Logs callable output. - This is useful for logging output of passes. Note inplace_wrapper replaces + This is useful for logging output of passes. Note ``inplace_wrapper`` replaces the pass output with the modified object. If we want to log the original - output, apply this wrapper before inplace_wrapper. + output, apply this wrapper before ``inplace_wrapper``. + Example:: - ``` - def my_pass(d: Dict) -> bool: - changed = False - if "foo" in d: - d["foo"] = "bar" - changed = True - return changed + def my_pass(d: Dict) -> bool: + changed = False + if "foo" in d: + d["foo"] = "bar" + changed = True + return changed - pm = PassManager(passes=[inplace_wrapper(log_hook(my_pass))]) - ``` + pm = PassManager(passes=[inplace_wrapper(log_hook(my_pass))]) Args: fn (Callable[Type1, Type2]) @@ -69,8 +75,8 @@ def my_pass(d: Dict) -> bool: """ @wraps(fn) - def wrapped_fn(gm): - val = fn(gm) + def wrapped_fn(*args: _P.args, **kwargs: _P.kwargs) -> _R: + val = fn(*args, **kwargs) logger.log(level, "Ran pass %s\t Return value: %s", fn, val) return val @@ -78,10 +84,10 @@ def wrapped_fn(gm): def loop_pass( - base_pass: Callable, + base_pass: Callable[[_T], _T], n_iter: int | None = None, - predicate: Callable | None = None, -): + predicate: Callable[[_T], bool] | None = None, +) -> Callable[[_T], _T]: """ Convenience wrapper for passes which need to be applied multiple times. @@ -97,7 +103,7 @@ def loop_pass( raise AssertionError("Exactly one of `n_iter`or `predicate` must be specified.") @wraps(base_pass) - def new_pass(source): + def new_pass(source: _T) -> _T: output = source if n_iter is not None and n_iter > 0: for _ in range(n_iter): @@ -120,8 +126,9 @@ def new_pass(source): # Implemented as 'depends on' operators. A constraint is satisfied iff a list # has a valid partial ordering according to this comparison operator. def _validate_pass_schedule_constraint( - constraint: Callable[[Callable, Callable], bool], passes: list[Callable] -): + constraint: Callable[[Callable[..., Any], Callable[..., Any]], bool], + passes: list[Callable[..., Any]], +) -> None: for i, a in enumerate(passes): for j, b in enumerate(passes[i + 1 :]): if constraint(a, b): @@ -133,42 +140,45 @@ def _validate_pass_schedule_constraint( ) -def this_before_that_pass_constraint(this: Callable, that: Callable): +def this_before_that_pass_constraint( + this: Callable[..., Any], that: Callable[..., Any] +) -> Callable[[Callable[..., Any], Callable[..., Any]], bool]: """ Defines a partial order ('depends on' function) where `this` must occur before `that`. """ - def depends_on(a: Callable, b: Callable): + def depends_on(a: Callable[..., Any], b: Callable[..., Any]) -> bool: return a != that or b != this return depends_on -def these_before_those_pass_constraint(these: Callable, those: Callable): +def these_before_those_pass_constraint( + these: Callable[..., Any], those: Callable[..., Any] +) -> Callable[[Callable[..., Any], Callable[..., Any]], bool]: """ - Defines a partial order ('depends on' function) where `these` must occur - before `those`. Where the inputs are 'unwrapped' before comparison. + Defines a partial order ('depends on' function) where ``these`` must occur + before ``those``. Where the inputs are 'unwrapped' before comparison. + + For example, the following pass list and constraint list would be invalid:: - For example, the following pass list and constraint list would be invalid. - ``` - passes = [ - loop_pass(pass_b, 3), - loop_pass(pass_a, 5), - ] + passes = [ + loop_pass(pass_b, 3), + loop_pass(pass_a, 5), + ] - constraints = [these_before_those_pass_constraint(pass_a, pass_b)] - ``` + constraints = [these_before_those_pass_constraint(pass_a, pass_b)] Args: these (Callable): pass which should occur first those (Callable): pass which should occur later Returns: - depends_on (Callable[[Object, Object], bool] + depends_on (Callable[[Object, Object], bool]) """ - def depends_on(a: Callable, b: Callable): + def depends_on(a: Callable[..., Any], b: Callable[..., Any]) -> bool: return unwrap(a) != those or unwrap(b) != these return depends_on @@ -190,40 +200,42 @@ class PassManager: `this_before_that_pass_constraint` for example. """ - passes: list[Callable] - constraints: list[Callable] + passes: list[Callable[..., Any]] + constraints: list[Callable[..., Any]] _validated: bool = False def __init__( self, - passes=None, - constraints=None, - ): + passes: list[Callable[..., Any]] | None = None, + constraints: list[Callable[..., Any]] | None = None, + ) -> None: self.passes = passes or [] self.constraints = constraints or [] @classmethod - def build_from_passlist(cls, passes): + def build_from_passlist(cls, passes: list[Callable[..., Any]]) -> "PassManager": pm = PassManager(passes) # TODO(alexbeloi): add constraint management/validation return pm - def add_pass(self, _pass: Callable): + def add_pass(self, _pass: Callable[..., Any]) -> None: self.passes.append(_pass) self._validated = False - def add_constraint(self, constraint): + def add_constraint(self, constraint: Callable[..., Any]) -> None: self.constraints.append(constraint) self._validated = False - def remove_pass(self, _passes: list[str]): + def remove_pass(self, _passes: list[str]) -> None: if _passes is None: return passes_left = [ps for ps in self.passes if ps.__name__ not in _passes] self.passes = passes_left self._validated = False - def replace_pass(self, _target, _replacement): + def replace_pass( + self, _target: Callable[..., Any], _replacement: Callable[..., Any] + ) -> None: passes_left = [] for ps in self.passes: if ps.__name__ == _target.__name__: @@ -233,7 +245,7 @@ def replace_pass(self, _target, _replacement): self.passes = passes_left self._validated = False - def validate(self): + def validate(self) -> None: """ Validates that current pass schedule defined by `self.passes` is valid according to all constraints in `self.constraints` @@ -244,7 +256,7 @@ def validate(self): _validate_pass_schedule_constraint(constraint, self.passes) self._validated = True - def __call__(self, source): + def __call__(self, source: Any) -> Any: self.validate() out = source for _pass in self.passes: diff --git a/torch/fx/passes/regional_inductor.py b/torch/fx/passes/regional_inductor.py index aa830c705ef72..42c11f1c19fbb 100644 --- a/torch/fx/passes/regional_inductor.py +++ b/torch/fx/passes/regional_inductor.py @@ -1,7 +1,11 @@ -# mypy: allow-untyped-defs - import functools import logging +from collections.abc import Callable, Mapping +from typing import Any, ParamSpec, TypeVar + + +_P = ParamSpec("_P") +_R = TypeVar("_R") import torch from torch.fx._compatibility import compatibility @@ -15,34 +19,15 @@ # standalone_inductor returns a callable class object - this does not sit well # with Fx graph node op call_function which expects a function. So this is just # a wrapper function to make Fx graph codegen happy. -def _dummy_wrapper(fn): +def _dummy_wrapper(fn: Callable[_P, _R]) -> Callable[_P, _R]: @functools.wraps(fn) - def inner(*args, **kwargs): + def inner(*args: _P.args, **kwargs: _P.kwargs) -> _R: return fn(*args, **kwargs) return inner -def _partition_by_supported_nodes(gm, supported_ops, prefix): - from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner - from torch.fx.passes.utils.fuser_utils import fuse_by_partitions - - partitioner = CapabilityBasedPartitioner( - gm, supported_ops, allows_single_node_partition=True - ) - - candidate_partitions = partitioner.propose_partitions() - partitioned_gm = fuse_by_partitions( - partitioner.graph_module, - [partition.nodes for partition in candidate_partitions], - prefix=prefix, - always_return_tuple=True, - ) - - return partitioned_gm - - -def _compile_submod(gm, prefix): +def _compile_submod(gm: torch.fx.GraphModule, prefix: str) -> torch.fx.GraphModule: from torch._inductor.standalone_compile import AOTCompiledArtifact for node in gm.graph.nodes: @@ -61,7 +46,7 @@ def _compile_submod(gm, prefix): # Get inductor configs from annotation # TODO we should change partition when there are multiple differently # annotated regions. - inductor_options = {} + inductor_options: dict[str, Any] = {} for sub_node in submod.graph.nodes: if hasattr(sub_node, "meta") and sub_node.meta.get("custom", None): custom = sub_node.meta["custom"] @@ -118,8 +103,8 @@ def _compile_submod(gm, prefix): return gm -def _needs_inductor_compile(node: torch.fx.Node): - return ( +def _needs_inductor_compile(node: torch.fx.Node) -> bool: + return bool( node.op not in ("placeholder", "output") and hasattr(node, "meta") and node.meta.get("custom", None) @@ -133,30 +118,66 @@ class _RegionScooper: """ @staticmethod - def scoop_regions(gm): - from torch.fx.passes.operator_support import OperatorSupport - - found_marked_node = False + def scoop_regions(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: + from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner + from torch.fx.passes.operator_support import create_op_support + from torch.fx.passes.utils.fuser_utils import fuse_by_partitions + + # Group tagged nodes by region ID. The region ID comes from the + # optional "inductor_region" key inside the compile_with_inductor + # annotation. When absent, all tagged nodes share a single default region + _DEFAULT_REGION = object() + regions: dict[object, set[torch.fx.Node]] = {} for node in gm.graph.nodes: if _needs_inductor_compile(node): - found_marked_node = True - break + compile_value = node.meta["custom"]["compile_with_inductor"] + if ( + isinstance(compile_value, dict) + and "inductor_region" in compile_value + ): + rid = compile_value["inductor_region"] + else: + rid = _DEFAULT_REGION + regions.setdefault(rid, set()).add(node) - if not found_marked_node: + if not regions: logger.info("No inductor marked nodes found") return gm - class InductorMarkedNodes(OperatorSupport): - def is_node_supported(self, submodules, node: torch.fx.Node) -> bool: - return _needs_inductor_compile(node) - - marked_nodes = InductorMarkedNodes() - return _partition_by_supported_nodes( - gm, marked_nodes, "__marked_inductor_submod" + # Run CapabilityBasedPartitioner per region to get cycle-safe partitions + # without merging across region boundaries. + def _is_in_region( + region_nodes: set[torch.fx.Node], + ) -> Callable[[Mapping[str, torch.nn.Module], torch.fx.Node], bool]: + def is_node_supported( + _submodules: Mapping[str, torch.nn.Module], node: torch.fx.Node + ) -> bool: + return node in region_nodes + + return is_node_supported + + all_partitions: list[dict[torch.fx.Node, int | None]] = [] + for region_nodes in regions.values(): + support = create_op_support(_is_in_region(region_nodes)) + partitioner = CapabilityBasedPartitioner( + gm, support, allows_single_node_partition=True + ) + for partition in partitioner.propose_partitions(): + all_partitions.append(partition.nodes) + + return fuse_by_partitions( + gm, + all_partitions, + prefix="__marked_inductor_submod", + always_return_tuple=True, ) @staticmethod - def recursively_scoop_regions(gm): + def recursively_scoop_regions( + gm: torch.fx.GraphModule, _processed: set[int] | None = None + ) -> torch.fx.GraphModule: + if _processed is None: + _processed = set() for node in gm.graph.find_nodes(op="get_attr"): if _needs_inductor_compile(node): # If the get_attr itself is marked for compile, the outer graph will @@ -164,12 +185,17 @@ def recursively_scoop_regions(gm): # regional inductor compiles that do not work well. continue submod = getattr(gm, node.target) - if isinstance(submod, torch.fx.GraphModule): - _RegionScooper.recursively_scoop_regions(submod) + # Track by id: multiple get_attr nodes may reference the same GraphModule + if ( + isinstance(submod, torch.fx.GraphModule) + and id(submod) not in _processed + ): + _processed.add(id(submod)) + _RegionScooper.recursively_scoop_regions(submod, _processed) return _RegionScooper.scoop_regions(gm) - def __call__(self, gm): + def __call__(self, gm: torch.fx.GraphModule) -> torch.fx.GraphModule: with torch.fx.traceback.preserve_node_meta(enable=False): return _RegionScooper.recursively_scoop_regions(gm) @@ -180,7 +206,7 @@ class _RegionCompiler: """ @staticmethod - def compile_region(gm): + def compile_region(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: from torch.fx.graph import _BoxedCodeGen gm = _compile_submod(gm, "__marked_inductor_submod") @@ -189,7 +215,7 @@ def compile_region(gm): return gm @staticmethod - def recursively_compile_regions(gm): + def recursively_compile_regions(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: # Find if the graph module has a scooped out region found_region = False for node in gm.graph.find_nodes(op="call_module"): @@ -208,35 +234,41 @@ def recursively_compile_regions(gm): return _RegionCompiler.compile_region(gm) return gm - def __call__(self, gm): + def __call__(self, gm: torch.fx.GraphModule) -> torch.fx.GraphModule: with torch.fx.traceback.preserve_node_meta(enable=False): return _RegionCompiler.recursively_compile_regions(gm) -def _create_inductor_marked_regions(gm): +def _create_inductor_marked_regions(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: with torch.fx.traceback.preserve_node_meta(enable=False): return _RegionScooper()(gm) -def _compile_inductor_marked_regions(gm): +def _compile_inductor_marked_regions(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: with torch.fx.traceback.preserve_node_meta(enable=False): return _RegionCompiler()(gm) @compatibility(is_backward_compatible=False) -def regional_inductor(gm, *example_args): +def regional_inductor( + gm: torch.fx.GraphModule, *example_args: object +) -> torch.fx.GraphModule: """ Scoops out inductor marked regions and compiles them with inductor. - Inductor options should be provided via the annotation API: - with fx_traceback.annotate({ - "compile_with_inductor": { - "inductor_configs": { - "max_autotune": True, - "triton.cudagraphs": False + Inductor options should be provided via the annotation API:: + + with fx_traceback.annotate( + { + "compile_with_inductor": { + "inductor_configs": { + "max_autotune": True, + "triton.cudagraphs": False, + } + } } - } - }): + ): + ... """ # fuser utils create new nodes using create_proxy which retains the seq_nr @@ -248,5 +280,5 @@ def regional_inductor(gm, *example_args): if torch._functorch.config.force_autograd_cache: from torch._inductor.output_code import RegionalOutputCode - gm = RegionalOutputCode(gm) + return RegionalOutputCode(gm) # type: ignore[return-value] return gm diff --git a/torch/fx/passes/regional_inductor_invoke_subgraph.py b/torch/fx/passes/regional_inductor_invoke_subgraph.py index b2a6e5c3a9b28..5d5f6d6e4d0db 100644 --- a/torch/fx/passes/regional_inductor_invoke_subgraph.py +++ b/torch/fx/passes/regional_inductor_invoke_subgraph.py @@ -1,5 +1,3 @@ -# mypy: allow-untyped-defs - import copy import logging from collections import defaultdict @@ -18,7 +16,7 @@ def _compile_submod( gm: torch.fx.GraphModule, subgraph: str, subgraph_users: list[torch.fx.Node] -): +) -> torch.fx.GraphModule: """ Compiles subgraph submodule in gm. subgraph is used by subgraph_users. subgraph_users must all be torch.ops.higher_order.invoke_subgraph HOP. @@ -56,7 +54,7 @@ def _compile_submod( compile_config, ) - def get_compiled_fn(): + def get_compiled_fn() -> AOTCompiledArtifact: context = torch._guards.TracingContext.get() if context.fake_mode is None: raise AssertionError("context.fake_mode is None") @@ -96,28 +94,33 @@ def get_compiled_fn(): return gm -def _needs_inductor_compile(node: torch.fx.Node): +def _needs_inductor_compile(node: torch.fx.Node) -> bool: # TODO: maybe we could change to check # node.meta.get("partitioner_tag") != "is_forward" # if the tag is relibable - return ( - node.op not in ("placeholder", "output") - and hasattr(node, "meta") - and node.meta.get("custom", None) - and node.meta["custom"].get("nested_region_config", None) - and node.meta["custom"]["nested_region_config"].fw_compiler - and node.meta.get("partitioner_tag") != "is_backward" - ) or ( - node.op not in ("placeholder", "output") - and hasattr(node, "meta") - and node.meta.get("custom", None) - and node.meta["custom"].get("nested_region_config", None) - and node.meta["custom"]["nested_region_config"].bw_compiler - and node.meta.get("partitioner_tag") == "is_backward" + return bool( + ( + node.op not in ("placeholder", "output") + and hasattr(node, "meta") + and node.meta.get("custom", None) + and node.meta["custom"].get("nested_region_config", None) + and node.meta["custom"]["nested_region_config"].fw_compiler + and node.meta.get("partitioner_tag") != "is_backward" + ) + or ( + node.op not in ("placeholder", "output") + and hasattr(node, "meta") + and node.meta.get("custom", None) + and node.meta["custom"].get("nested_region_config", None) + and node.meta["custom"]["nested_region_config"].bw_compiler + and node.meta.get("partitioner_tag") == "is_backward" + ) ) -def _compile_invoke_subgraph_nodes_with_inductor(gm): +def _compile_invoke_subgraph_nodes_with_inductor( + gm: torch.fx.GraphModule, +) -> torch.fx.GraphModule: map_subgraph_to_nodes = defaultdict(list) subgraphs: set[str] = set() @@ -140,7 +143,9 @@ def _compile_invoke_subgraph_nodes_with_inductor(gm): return gm -def _recursive_compile_invoke_subgraph_nodes(gm): +def _recursive_compile_invoke_subgraph_nodes( + gm: torch.fx.GraphModule, +) -> torch.fx.GraphModule: for node in gm.graph.find_nodes(op="get_attr"): if _needs_inductor_compile(node): # If the get_attr itself is marked for compile, the outer graph will @@ -155,7 +160,9 @@ def _recursive_compile_invoke_subgraph_nodes(gm): @compatibility(is_backward_compatible=False) -def regional_inductor_invoke_subgraph(gm, *example_args): +def regional_inductor_invoke_subgraph( + gm: torch.fx.GraphModule, *example_args: object +) -> torch.fx.GraphModule: """ Compile invoke_subgraph nodes if they have custom compiler specified in node.meta["nested_region_config"].bw_compiler or fw_compiler @@ -165,6 +172,7 @@ def regional_inductor_invoke_subgraph(gm, *example_args): with torch.fx.traceback.preserve_node_meta(enable=False): compiled_gm = _recursive_compile_invoke_subgraph_nodes(gm) # TODO: might not need this boxed_nop after we switch to _RegionCompiler + # pyrefly: ignore [bad-return] return torch._dynamo.backends.debugging.boxed_nop( compiled_gm, example_inputs=[] ) diff --git a/torch/fx/passes/reinplace.py b/torch/fx/passes/reinplace.py index 17dde9256f5fa..0875e9ace0ea7 100644 --- a/torch/fx/passes/reinplace.py +++ b/torch/fx/passes/reinplace.py @@ -1,8 +1,6 @@ -# mypy: allow-untyped-defs import _operator import itertools from collections import defaultdict -from collections.abc import Callable from enum import Enum from typing import Any @@ -24,7 +22,7 @@ class _ViewType(Enum): MultiOutputView = 2 -def _is_view_op(tgt): +def _is_view_op(tgt: object) -> bool | None: if tgt is not None and isinstance(tgt, torch._ops.OpOverload): schema = tgt._schema if len(schema.arguments) > 0: @@ -35,7 +33,7 @@ def _is_view_op(tgt): ) -def _get_view_type(tgt) -> _ViewType: +def _get_view_type(tgt: object) -> _ViewType: if tgt is not None and isinstance(tgt, torch._ops.OpOverload): schema = tgt._schema if len(schema.arguments) > 0: @@ -61,7 +59,7 @@ def _get_view_type(tgt) -> _ViewType: # to sanity check that our aliasing information is correct. @compatibility(is_backward_compatible=False) class _FunctionalizationMetadataProp(torch.fx.Interpreter): - def run_node(self, node: Node): + def run_node(self, node: Node) -> Any: self.node_counter += 1 result = super().run_node(node) node.meta["fake_result"] = result @@ -122,7 +120,7 @@ def run_node(self, node: Node): raise AssertionError("view_storage != base_storage") return result - def propagate(self, *args): + def propagate(self, *args: object) -> object: self.multi_output_view_nodes = {} self.node_counter = -1 @@ -133,7 +131,9 @@ def propagate(self, *args): return super().run(*fake_args) -def _schemas_match(functional_schema, inplace_schema): +def _schemas_match( + functional_schema: torch._C.FunctionSchema, inplace_schema: torch._C.FunctionSchema +) -> bool: names_match = ( inplace_schema.name.endswith("_") and inplace_schema.name[:-1] == functional_schema.name @@ -160,7 +160,7 @@ def _schemas_match(functional_schema, inplace_schema): # - mutating ops (e.g. _fused_moving_avg_obs_fq_helper) # - out= ops (e.g. angle -> angle.out) # TODO: we should also figure this info out using torchgen. -def _maybe_get_inplace_op(op): +def _maybe_get_inplace_op(op: object) -> torch._ops.OpOverload | None: # __module__ seems broken; it returns torch._ops.aten which doesn't exist if not isinstance(op, torch._ops.OpOverload): return None @@ -202,7 +202,7 @@ def _maybe_get_inplace_op(op): return inplace_op -_VIEW_INVERSE_MAP: dict[Callable[..., Any], Callable[..., Any]] = { +_VIEW_INVERSE_MAP: dict[torch._ops.OpOverload, torch._ops.OpOverload] = { torch.ops.aten.diagonal_scatter.default: torch.ops.aten.diagonal.default, torch.ops.aten.select_scatter.default: torch.ops.aten.select.int, torch.ops.aten.slice_scatter.default: torch.ops.aten.slice.Tensor, @@ -213,8 +213,8 @@ def _maybe_get_inplace_op(op): # This function, given a set of set of (aliased) tensor nodes, # Returns any nodes in the graph that *use* any of the aliases, that occur *after* op_index # in the node ordering. -def _get_all_later_node_usages(tensor_aliases: set[Node], op_index: int): - def _add_if_tensor(x, set_): +def _get_all_later_node_usages(tensor_aliases: set[Node], op_index: int) -> set[Node]: + def _add_if_tensor(x: object, set_: set[StorageWeakRef]) -> None: if isinstance(x, FakeTensor): set_.add(StorageWeakRef(x._typed_storage())) @@ -249,7 +249,7 @@ def _add_if_tensor(x, set_): def _get_view_inverse_node_usages( later_node_usages: set[Node], self_aliases: set[Node] ) -> set[Node]: - def matching_view_metadata(a, b): + def matching_view_metadata(a: FakeTensor, b: FakeTensor) -> bool: return ( a.size() == b.size() and a.stride() == b.stride() @@ -308,105 +308,111 @@ def matching_view_metadata(a, b): @compatibility(is_backward_compatible=True) -def reinplace(gm, *sample_args): - """ +def reinplace( + gm, *sample_args +): # pyrefly: ignore[unannotated-parameter, unannotated-return] + r""" Given an fx.GraphModule, modifies it to perform "reinplacing", mutating the nodes of the graph. - We look for out-of-place op call sites like `b = a.add(...)`, - and convert them to be inplace (`b = a.add_(...)`), + We look for out-of-place op call sites like ``b = a.add(...)``, + and convert them to be inplace (``b = a.add_(...)``), as long as the input to the current operator ("a") isn't reused anywhere later in the graph. This pass currently expects to operate on a **functional, ATen** graph. - This can be obtained by running `make_fx(functionalize(f))`. + This can be obtained by running ``make_fx(functionalize(f))``. Sample inputs are needed to determine aliasing relationships of the inputs. - In general, we can't reinplace node `b = a.add(...)` if "a" aliases any of the + In general, we can't reinplace node ``b = a.add(...)`` if "a" aliases any of the inputs to the program. - Given a node "b = foo(a, args...) the algorithm for re-inplacing is as follows: + Given a node ``b = foo(a, args...)`` the algorithm for re-inplacing is as follows: + + **(1)** Perform some initial checks on the metadata of "a" and "args..." + that can disqualify them from being reinplaced. + + - **(1a)** Check that the self argument we're attempting to reinplace + has acceptable dtype/size metadata to reinplace with. + + For example, if we have:: + + a = torch.ones(1) + b = torch.ones(10) + out = torch.add(a, b) + + We can't turn that into ``a.add_(b)`` because that would require resizing "a". - (1) Perform some initial checks on the metadata of "a" and "args..." - that can disqualify them from being reinplaced. + Similarly, we can't convert ``torch.ge(a, b)`` into ``a.ge_(b)``, + because that would require changing a's dtype (from e.g. float32 to bool). + Note that in this specific example, we could technically do better.. - (1a) Check that the self argument we're attempting to reinplace - has acceptable dtype/size metadata to reinplace with. + If we see the pattern:: - For example, if we have: - a = torch.ones(1) - b = torch.ones(10) - out = torch.add(a, b) - We can't turn that into - a.add_(b) - Because that would require resizing "a". + a_1 = a.ge(b) + a_2 = aten._to_copy(a_1, a.dtype) - Similarly, we can't convert torch.ge(a, b) into a.ge_(b), - because that would require changing a's dtype (from e.g. float32 to bool). - Note that in this specific example, we could technically do better.. + Then this should be valid to completely re-inplace + (this is exactly what functionalization will emit when it sees ``a.ge_(b)``). - If we see the pattern: - a_1 = a.ge(b) - a_2 = aten._to_copy(a_1, a.dtype) - Then we this should be valid to completely re-inplace - (this is exactly what functionalization will emit when it sees a.ge_(b)). + This optimization is only really important for user programs + that directly use inplace comparison ops though. - This optimization is only really important for user programs - that directly use inplace comparison ops though. + We also cannot re-inplace on tensors that have overlapping memory, + e.g. ``torch.ones(1).expand(4, 4).add_(1)``. - We also cannot re-inplace on tensors that have overlapping memory, - e.g. torch.ones(1).expand(4, 4).add_(1) + - **(1b)** Check if "a" is an alias of any of the program inputs. - (1b) Check if "a" is an alias of any of the program inputs. + If it is, skip and move to the next node. + Inplace'ing an op that would cause it to mutate a program is not sound, + because that would be a side effect visible to the user. - If it is, skip and move to the next node. - Inplace'ing an op that would cause it to mutate a program is not sound, - because that would be a side effect visible to the user. + NOTE: there's a future optimization that we should make: + if "a" is a (alias of a) program input, but later in the program + there is a node that looks like ``a.copy_(...)``, + then re-inplacing is ok to do - we are temporarily reusing a's buffer, + which will later be overwritten by the ``copy_()`` call. - NOTE: there's a future optimization that we should make: - if "a" is a (alias of a) program input, but later in the program - there is a node that looks like "a.copy_(...)", - Then re-inplacing is ok to do - we are temporarily reusing a's buffer, - which will later be overwritten by the copy_() call. + This will be an important optimization to have for programs that mutate + their inputs. It currently isn't implemented though. - This will be an important optimization to have for programs that mutate - their inputs. It currently isn't implemented though. + - **(1c)** Check if "a" and "args..." alias. - (1c) Check if "a" and "args..." alias + For example, re-inplacing to create code like the below + isn't guaranteed to be sound:: - For example, re-inplacing to create code like the below - isn't guaranteed to be sound: + aten.mul_(a, a) - aten.mul_(a, a) + **(2)** Check that "a" and all of its outstanding aliases are not used anywhere + later in the graph. If this is the case, then it's safe to re-inplace + to ``b = foo_(a)``. - (2) Check that "a" and all of its outstanding aliases are not used anywhere - later in the graph. If this is the case, then it's safe to re-inplace - to "b = foo_(a)". + There are a few caveats to this, explained in more detail below: - There are a few caveats to this, explained in more detail below: - (a) If "a" is used later as an argument to a view op, that is okay. - It's only a problem if "a" (or that view) is later passed - into a normal operator, or if it is returned as the program output. - (b) If "a" is a repeat argument in `foo()`, then don't reinplace. - Most ATen kernels don't make any guarantees that this is sound, - e.g. if you do aten.mul_(a, a). - So we'll just ban re-inplacing in this case. - It's only a problem if "a" (or that view) is later passed - (c) If "a" is used as an input into a view "inverse" / "scatter" - operator, it is potentially fine to re-inplace - (and remove that scatter operator from the graph). - See below for a more detailed example. + - (a) If "a" is used later as an argument to a view op, that is okay. + It's only a problem if "a" (or that view) is later passed + into a normal operator, or if it is returned as the program output. + - (b) If "a" is a repeat argument in ``foo()``, then don't reinplace. + Most ATen kernels don't make any guarantees that this is sound, + e.g. if you do ``aten.mul_(a, a)``. + So we'll just ban re-inplacing in this case. + - (c) If "a" is used as an input into a view "inverse" / "scatter" + operator, it is potentially fine to re-inplace + (and remove that scatter operator from the graph). + See below for a more detailed example. - NOTE: there is an optimization in this step that is crucial - to fully recovering performance from functionalization. + NOTE: there is an optimization in this step that is crucial + to fully recovering performance from functionalization. + + Given this program:: - Given this program: def f(x): a = torch.ops.aten.add(x, x) b = torch.ops.aten.diagonal(a) torch.ops.aten.fill_(b, 0) return d - Functionalization will emit the following: + Functionalization will emit the following:: + def f(x): a = torch.ops.aten.add(x, x) b = torch.ops.aten.diagonal(a, 0, 1) @@ -414,99 +420,115 @@ def f(x): a_updated = torch.ops.aten.diagonal_scatter(a, b_updated, 0, 1) return a_updated - Ordinarily, we would not be able to reinplace the fill, - because "b" aliases with "a" which is used by the diagonal_scatter call. - - "re-inplacing" is on the hook for figuring out that it is ok to - completely, the expensive diagonal_scatter call, if we re-inplace the add(). - - So, for every `alias in alias_set(a)`, instead of checking - that "alias" is not used anywhere later in the graph, - we check that - EITHER: - (a) alias is not used anywhere later in the graph - OR: - (b) alias is used exactly once later on in the graph, - in the following op: - - out = foo_scatter(alias, x, args...) - - where the following must hold: - (i) "foo_scatter" is the "inverse" operator for foo. - This only applies to "foo" ops that are view operators, - which view into a subset of the original tensor's memory. - In practice, there are ~4 operators where this applies: - diagonal -> diagonal_scatter - slice -> slice_scatter - select -> select_scatter - as_strided -> as_strided_scatter - (ii) "args..." are the same between the foo() and foo_scatter() calls. - - (3) Perform the actual re-inplacing on foo! - - (3b) is the common case, but special care is needed for {view}_scatter (3a) - - (3a) {view}_scatter ops. - - Consider this program: - a = torch.zeros(2, 2) - b = torch.ones(2) - a[0] = b - - Post functionalization, that will look like: - a = torch.zeros(2) - b = torch.ones(1) - a_updated = torch.select_scatter(a, b, 0, 0) - - In this case though, there is no "functional" op to re-inplace! - Instead, we'd like to directly remove toe select_scatter call. - We already know from (3) that this is valid, - because "a" has no later usages in the graph. - - We perform the re-inplacing on the {view}_scatter op like so - Before: - a_updated = torch.select_scatter(a, b, args...) - After: - a_slice = a.select(a, args...) - a_slice.copy_(b) - - (3b) Otherwise, replace the functional op with its inplace variant. - Before: - b = foo(a, args...) - After: - a.foo_(args...) - - (4) Finally, after converting either: - Before: - b = foo(a) - After: - foo_(a) - or - Before: - b = {slice}_scatter(a, mutated_slice, args...) - After: - slice = {slice}(a, args...) - slice.copy_(mutated_slice) - - We now need to find all later nodes that use "b" as an argument - and update them to take in "a" instead. - - Note that for the majority of inplace ops, this isn't actually necessary - (because most inplace ops return "self" as their output). - This isn't generally true for all mutable ops though, which is why - we need to actually replace all of the arguments. - - We also need to update our metadata of Dict[StorageWeakRef, Set[Node]], - That maps a given tensor storage to the set of all nodes that take in that storage - as an input. - Specifically, re-inplacing `b = foo(a)` causes "a" and "b"'s sets to get fused - together. - - (5) Any "view_inverse/scatter" nodes that were identified as "it's ok to ignore them" - during step (3) get manually deleted from the graph. - Their outputs are no longer used, so technically standard DCE would be able - to do this, but we can no longer run FX's DCE pass now that we have mutable - ops in the graph. + Ordinarily, we would not be able to reinplace the fill, + because "b" aliases with "a" which is used by the diagonal_scatter call. + + "re-inplacing" is on the hook for figuring out that it is ok to + completely remove the expensive diagonal_scatter call, if we re-inplace + the add(). + + So, for every ``alias in alias_set(a)``, instead of checking + that "alias" is not used anywhere later in the graph, + we check that EITHER: + + - (a) alias is not used anywhere later in the graph, OR + - (b) alias is used exactly once later on in the graph, + in the following op:: + + out = foo_scatter(alias, x, args...) + + where the following must hold: + + - (i) ``foo_scatter`` is the "inverse" operator for foo. + This only applies to "foo" ops that are view operators, + which view into a subset of the original tensor's memory. + In practice, there are ~4 operators where this applies:: + + diagonal -> diagonal_scatter + slice -> slice_scatter + select -> select_scatter + as_strided -> as_strided_scatter + + - (ii) "args..." are the same between the ``foo()`` and + ``foo_scatter()`` calls. + + **(3)** Perform the actual re-inplacing on foo! + + (3b) is the common case, but special care is needed for + ``{view}_scatter`` (3a). + + - **(3a)** ``{view}_scatter`` ops. + + Consider this program:: + + a = torch.zeros(2, 2) + b = torch.ones(2) + a[0] = b + + Post functionalization, that will look like:: + + a = torch.zeros(2) + b = torch.ones(1) + a_updated = torch.select_scatter(a, b, 0, 0) + + In this case though, there is no "functional" op to re-inplace! + Instead, we'd like to directly remove the select_scatter call. + We already know from (3) that this is valid, + because "a" has no later usages in the graph. + + We perform the re-inplacing on the ``{view}_scatter`` op like so. + + Before:: + + a_updated = torch.select_scatter(a, b, args...) + + After:: + + a_slice = a.select(a, args...) + a_slice.copy_(b) + + - **(3b)** Otherwise, replace the functional op with its inplace variant. + + Before:: + + b = foo(a, args...) + + After:: + + a.foo_(args...) + + **(4)** Finally, after converting either:: + + # Before: # After: + b = foo(a) foo_(a) + + or:: + + # Before: + b = {slice}_scatter(a, mutated_slice, args...) + # After: + slice = {slice}(a, args...) + slice.copy_(mutated_slice) + + We now need to find all later nodes that use "b" as an argument + and update them to take in "a" instead. + + Note that for the majority of inplace ops, this isn't actually necessary + (because most inplace ops return "self" as their output). + This isn't generally true for all mutable ops though, which is why + we need to actually replace all of the arguments. + + We also need to update our metadata of ``Dict[StorageWeakRef, Set[Node]]``, + that maps a given tensor storage to the set of all nodes that take in that + storage as an input. + Specifically, re-inplacing ``b = foo(a)`` causes "a" and "b"'s sets to get + fused together. + + **(5)** Any ``view_inverse/scatter`` nodes that were identified as + "it's ok to ignore them" during step (3) get manually deleted from the graph. + Their outputs are no longer used, so technically standard DCE would be able + to do this, but we can no longer run FX's DCE pass now that we have mutable + ops in the graph. """ _FunctionalizationMetadataProp(gm).propagate(*sample_args) @@ -543,7 +565,7 @@ def f(x): for n in gm.graph.nodes: if "fake_result" in n.meta: # Tree-mapping because some ops can return lists of tensors. - def _add_to_map(x): + def _add_to_map(x: object) -> None: if isinstance(x, FakeTensor): storage_to_nodes[StorageWeakRef(x._typed_storage())].add(n) @@ -701,12 +723,14 @@ def _add_to_map(x): # We need to replace any later usages of "b" with "a" for old in itertools.chain([node], later_view_inverse_node_usages): new = old.args[0] + if not isinstance(new, Node): + raise AssertionError(f"Expected Node, got {type(new)}") nodes_to_update = [ n for n in old.users if n.meta["node_idx"] > node.meta["node_idx"] ] for node_to_update in nodes_to_update: - def replace_arg(a): + def replace_arg(a: Node) -> Node: if a == old: return new return a @@ -751,6 +775,7 @@ def replace_arg(a): and len(node_res_storage) == 1 and old_res_storage == node_res_storage ): + # pyrefly: ignore [missing-attribute] new_flattened_res = pytree.tree_leaves(new.meta["fake_result"]) new_res_storage = { StorageWeakRef(x._typed_storage()) diff --git a/torch/fx/passes/runtime_assert.py b/torch/fx/passes/runtime_assert.py index 99665d407e65d..9a6e3881e8afa 100644 --- a/torch/fx/passes/runtime_assert.py +++ b/torch/fx/passes/runtime_assert.py @@ -1,8 +1,8 @@ -# mypy: allow-untyped-defs import functools import logging import operator import sys +from collections.abc import Callable from typing import Any, Optional, TYPE_CHECKING @@ -98,12 +98,12 @@ def insert_deferred_runtime_asserts( _get_placeholder_expr, _has_uninterpretable_sympy_function, CallMethodKey, - cast_symbool_to_symint_guardless, ConvertIntKey, DivideByKey, free_symbols, InnerTensorKey, resolve_unbacked_bindings, + RuntimeAssert, ) from torch.utils._sympy.numbers import int_oo from torch.utils._sympy.reference import ( @@ -207,7 +207,9 @@ def _node_metadata_hook( Analysis = PythonReferenceAnalysis if export else OptimizedPythonReferenceAnalysis - def _sympy_interp(expr_to_proxy, expr): + def _sympy_interp( + expr_to_proxy: dict[sympy.Expr, fx.Proxy], expr: sympy.Expr + ) -> fx.Proxy: # sympy_interp() with hash consing from sympy import Integer, Number, Symbol from sympy.logic.boolalg import BooleanAtom @@ -240,7 +242,7 @@ def _is_bound_expr_for_symbol(expr: "sympy.Expr") -> bool: isinstance(rhs, sympy.Symbol) and isinstance(lhs, sympy.Number) ) - def add_runtime_asserts(ras): + def add_runtime_asserts(ras: list[RuntimeAssert]) -> None: for ra in ras: if ( # redundant @@ -310,7 +312,7 @@ def add_runtime_asserts(ras): and (example_value := _get_example_value(node)) is not None ): - def match_symbol(symint, cb): + def match_symbol(symint: object, cb: Callable[[], fx.Node]) -> None: if ( isinstance(symint, torch.SymInt) and isinstance(symint.node, SymNode) @@ -391,7 +393,7 @@ def match_symbol(symint, cb): and (sym_expr := _get_sym_val(node)) is not None ): # this guards against deleting calls like item() that produce new untracked symbols - def has_new_untracked_symbols(): + def has_new_untracked_symbols() -> bool: # pyrefly: ignore [missing-attribute] for symbol in sym_expr.free_symbols: if symbol not in expr_to_proxy: @@ -405,7 +407,7 @@ def has_new_untracked_symbols(): shape_env, node.meta.get("unbacked_bindings", {}) ) - def has_new_unbacked_bindings(): + def has_new_unbacked_bindings() -> bool: if resolved_unbacked_bindings is None: raise AssertionError("resolved_unbacked_bindings is None") for key in resolved_unbacked_bindings: @@ -482,7 +484,7 @@ def has_new_unbacked_bindings(): # TODO: some CSE when generating these nodes can probably # help reduce graph size and improve compile time - def go(node, keypath): + def go(node: fx.Node, keypath: tuple[object, ...]) -> fx.Node: if keypath == (): return node if ( @@ -535,9 +537,7 @@ def go(node, keypath): ) elif isinstance(keypath[0], ConvertIntKey): return go( - graph.call_function( - cast_symbool_to_symint_guardless, (node,) - ), + graph.call_function(torch.sym_ite, (node, 1, 0)), keypath[1:], ) elif isinstance(keypath[0], DivideByKey): @@ -629,7 +629,7 @@ def go(node, keypath): # assert and also explicitly refine the range # (refinement should not be necessary once runtime # asserts cause refinement, but that's NYI) - def convert(s): + def convert(s: Any) -> int | None: if s in (int_oo, -int_oo): return None try: @@ -637,44 +637,37 @@ def convert(s): except TypeError: return None - if ( - expr_to_proxy[i0].node.target - is not cast_symbool_to_symint_guardless + with _set_node_metadata_hook( + gm, + functools.partial( + _node_metadata_hook, + stack_trace=node.meta.get("stack_trace"), + nn_module_stack=node.meta.get("nn_module_stack"), + # nodes added in `apply_runtime_assertion_pass` will have the same annotation + # as the input node to the assertion + custom=node.meta.get("custom"), + ), ): - # TODO(pianpwk): calling sym_constrain_range_for_size or adding bound asserts - # raises AOTAutograd errors on cast_symbool_to_symint_guardless - - with _set_node_metadata_hook( - gm, - functools.partial( - _node_metadata_hook, - stack_trace=node.meta.get("stack_trace"), - nn_module_stack=node.meta.get("nn_module_stack"), - # nodes added in `apply_runtime_assertion_pass` will have the same annotation - # as the input node to the assertion - custom=node.meta.get("custom"), - ), - ): - if (min_val := convert(vr.lower)) is not None: - ge = _sympy_interp(expr_to_proxy, i0 >= min_val).node - graph.call_function( - torch.ops.aten._assert_scalar.default, - ( - ge, - f"Runtime assertion failed for expression {i0 >= min_val} on node '{ge}'", - ), - ) - added_asserts.add(i0 >= min_val) - if (max_val := convert(vr.upper)) is not None: - le = _sympy_interp(expr_to_proxy, i0 <= max_val).node - graph.call_function( - torch.ops.aten._assert_scalar.default, - ( - le, - f"Runtime assertion failed for expression {i0 <= max_val} on node '{le}'", - ), - ) - added_asserts.add(i0 <= max_val) + if (min_val := convert(vr.lower)) is not None: + ge = _sympy_interp(expr_to_proxy, i0 >= min_val).node + graph.call_function( + torch.ops.aten._assert_scalar.default, + ( + ge, + f"Runtime assertion failed for expression {i0 >= min_val} on node '{ge}'", + ), + ) + added_asserts.add(i0 >= min_val) + if (max_val := convert(vr.upper)) is not None: + le = _sympy_interp(expr_to_proxy, i0 <= max_val).node + graph.call_function( + torch.ops.aten._assert_scalar.default, + ( + le, + f"Runtime assertion failed for expression {i0 <= max_val} on node '{le}'", + ), + ) + added_asserts.add(i0 <= max_val) constrained_unbacked_symbols.add(i0) add_runtime_asserts(ras) diff --git a/torch/fx/passes/shape_prop.py b/torch/fx/passes/shape_prop.py index 63e35cae175cd..f6887e2d67c5c 100644 --- a/torch/fx/passes/shape_prop.py +++ b/torch/fx/passes/shape_prop.py @@ -1,5 +1,3 @@ -# mypy: ignore-errors - import traceback from typing import Any, NamedTuple diff --git a/torch/fx/passes/split_module.py b/torch/fx/passes/split_module.py index bf3a89cbfe1a4..f703748c8de5d 100644 --- a/torch/fx/passes/split_module.py +++ b/torch/fx/passes/split_module.py @@ -1,4 +1,3 @@ -# mypy: allow-untyped-defs import inspect import logging from collections import OrderedDict @@ -19,7 +18,7 @@ @compatibility(is_backward_compatible=True) class Partition: - def __init__(self, name: str): + def __init__(self, name: str) -> None: self.name: str = name self.submod_name = f"submod_{name}" self.node_names: list[str] = [] @@ -63,7 +62,8 @@ def split_module( keep_original_input_name: bool = True, *, partition_affix: str | None = None, -): + tuple_return: bool = False, +) -> GraphModule: """ Creates subgraphs out of main graph @@ -87,6 +87,9 @@ def split_module( have the same input names as the original graph. partition_affix: Optional[str]: If specified, the submodules' names will contain the affix, e.g. "submod__". + tuple_return: bool: If True, submodule outputs are always wrapped in a tuple, + even when there is only a single output value. This makes all subgraphs + conform to the convention expected by ``torch._inductor.compile_fx``. Returns: GraphModule: the module after split. @@ -173,7 +176,7 @@ def construct_graph( node: Node, base_mod_env: dict[str, Node], base_mod_attrs: dict[str, torch.fx.graph_module.GraphModule], - ): + ) -> tuple[dict[str, Node], dict[str, torch.fx.graph_module.GraphModule]]: if node.op == "placeholder": default_value = ( node.args[0] if len(node.args) > 0 else inspect.Signature.empty @@ -210,7 +213,7 @@ def construct_graph( orig_nodes: dict[str, Node] = {} symbol_to_node: dict[sympy.Symbol, Node] = {} - def record_cross_partition_use(def_node: Node, use_node: Node | None): + def record_cross_partition_use(def_node: Node, use_node: Node | None) -> None: from torch.fx.experimental.symbolic_shapes import free_symbols defined = getattr(def_node, "_fx_partition", None) @@ -259,7 +262,7 @@ def record_cross_partition_use(def_node: Node, use_node: Node | None): if defined is not None: use_partition.dependencies.setdefault(defined) - def instantiate_node_partition_mapping(node): + def instantiate_node_partition_mapping(node: Node) -> None: partition_idx = split_callback(node) partition_name = str(partition_idx) if partition_affix is not None: @@ -463,7 +466,7 @@ def instantiate_node_partition_mapping(node): # We don't pass in get_attr nodes as inputs to the partition, but # instead set them as targets and use getattr within the module - def add_placeholder(): + def add_placeholder() -> Node: if keep_original_input_name: name = inp else: @@ -604,15 +607,10 @@ def add_placeholder(): partition.environment[orig_nodes[name]] for name in partition.outputs ) - # skip output node generation if there are no output values - num_output_vals = len(output_vals) - if num_output_vals == 1: + if len(output_vals) == 1 and not tuple_return: partition.graph.output(output_vals[0]) - elif num_output_vals > 1: - partition.graph.output(output_vals) else: - # Invariant - Graph should always have an output node. - partition.graph.output(()) + partition.graph.output(output_vals) if keep_original_order: # first get the attr nodes required by this partition @@ -650,8 +648,8 @@ def add_placeholder(): ) num_outputs = len(partition.outputs) - if num_outputs > 1: - # Unpack multiple return values from submodule + if num_outputs > 1 or (num_outputs == 1 and tuple_return): + # Unpack return values from submodule output_val_proxy = torch.fx.proxy.Proxy(output_val) for i, output_name in enumerate(partition.outputs): base_mod_env[output_name] = output_val_proxy[i].node # type: ignore[index] diff --git a/torch/fx/passes/split_utils.py b/torch/fx/passes/split_utils.py index 08a4c5a685ec7..a81a5c0689884 100644 --- a/torch/fx/passes/split_utils.py +++ b/torch/fx/passes/split_utils.py @@ -1,6 +1,6 @@ -# mypy: allow-untyped-defs import copy from dataclasses import dataclass, field +from typing import TYPE_CHECKING import torch.fx from torch.fx._compatibility import compatibility @@ -10,6 +10,10 @@ from .tools_common import CALLABLE_NODE_OPS, is_node_output_tensor, NodeList +if TYPE_CHECKING: + from .splitter_base import Subgraph + + __all__ = [ "getattr_recursive", "setattr_recursive", @@ -20,7 +24,7 @@ @compatibility(is_backward_compatible=False) -def getattr_recursive(obj, name): +def getattr_recursive(obj: object, name: str) -> object: for layer in name.split("."): if isinstance(obj, torch.nn.ModuleList): if hasattr(obj, "_modules") and layer in obj._modules: @@ -35,7 +39,7 @@ def getattr_recursive(obj, name): @compatibility(is_backward_compatible=False) -def setattr_recursive(obj, attr, value): +def setattr_recursive(obj: object, attr: str, value: object) -> None: if "." not in attr: setattr(obj, attr, value) else: @@ -55,13 +59,13 @@ class Component: name: str # Stores the placeholder nodes in `graph`. - input_placeholders: list = field(default_factory=list) + input_placeholders: list[torch.fx.Node] = field(default_factory=list) # Store the nodes in original graph that are placeholder in `graph`. - orig_inputs: list = field(default_factory=list) + orig_inputs: list[torch.fx.Node] = field(default_factory=list) # Store the nodes in original graph that are outputs in `graph`. - orig_outputs: list = field(default_factory=list) + orig_outputs: list[torch.fx.Node] = field(default_factory=list) # Mapping from get_attr node in original graph to get_attr node in `graph`. getattr_maps: dict[torch.fx.Node, torch.fx.Node] = field(default_factory=dict) @@ -217,10 +221,14 @@ def flatten(x: torch.fx.node.Argument) -> NodeList: ) # Map a input of `node` to nodes in the component's graph. - def remap_func(x): + def remap_func(x: torch.fx.Node) -> torch.fx.Node: # If input is a get_attr node, copy it to current component's graph. # Returns the get_attr node in current component's graph. if x.op == "get_attr": + if not isinstance(x.target, str): + raise RuntimeError( + f"Expected get_attr node target to be a str, got {type(x.target)}" + ) if x not in comp.getattr_maps: comp.getattr_maps[x] = comp.graph.get_attr( x.target, type_expr=x.type @@ -321,7 +329,7 @@ def remap_func(x): @compatibility(is_backward_compatible=False) -def move_non_tensor_nodes_on_boundary(subgraphs) -> None: +def move_non_tensor_nodes_on_boundary(subgraphs: list["Subgraph"]) -> None: """ Move non-tensor nodes on the boundary between subgraphs. @@ -391,7 +399,7 @@ def can_move_node_and_dependencies( visited = set() can_move = True - def dfs(current_node): + def dfs(current_node: torch.fx.Node) -> None: nonlocal can_move, nodes_to_move if current_node in visited: diff --git a/torch/fx/passes/splitter_base.py b/torch/fx/passes/splitter_base.py index 3de6753228a82..0536c67ca9f98 100644 --- a/torch/fx/passes/splitter_base.py +++ b/torch/fx/passes/splitter_base.py @@ -1,12 +1,11 @@ -# mypy: allow-untyped-defs import argparse import copy import json import os from collections import defaultdict -from collections.abc import Iterable, Sequence +from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass -from typing import Any, Literal, NamedTuple +from typing import Any, IO, Literal, NamedTuple import torch from torch._logging import trace_structured @@ -76,12 +75,12 @@ class _SplitterSettingBase: def __init__( self, - min_acc_module_size=DEFAULT_MIN_ACC_MODULE_SIZE, - skip_fusion=DEFAULT_SKIP_FUSION, - allow_non_tensor=DEFAULT_ALLOW_NON_TENSOR, + min_acc_module_size: int = DEFAULT_MIN_ACC_MODULE_SIZE, + skip_fusion: bool = DEFAULT_SKIP_FUSION, + allow_non_tensor: bool = DEFAULT_ALLOW_NON_TENSOR, max_acc_splits: int = -1, move_non_tensor_nodes_on_boundary: bool = False, - ): + ) -> None: parser = argparse.ArgumentParser() parser.add_argument( "--min-acc-module-size", @@ -159,12 +158,12 @@ class NodeEvent: def __init__( self, source: torch.fx.Node, desc: str, dep: torch.fx.Node | None = None - ): + ) -> None: self.source = source self.desc = desc self.dep = dep - def to_str(self): + def to_str(self) -> str: # source: The name of the subject of the event. # desc: description of the event, in the format of | # dep: The name of the cause of this event, which is another node, or # @@ -178,7 +177,7 @@ class NodeEventTracker: Tracks node events during the splitter execution. """ - def __init__(self, tracker_mode, dump_prefix): + def __init__(self, tracker_mode: int, dump_prefix: str) -> None: self.tracker_mode = tracker_mode self.dump_prefix = dump_prefix # list of events @@ -187,7 +186,9 @@ def __init__(self, tracker_mode, dump_prefix): self.node_events = {} self.writer = print - def add(self, node: torch.fx.Node, desc: str, dep: torch.fx.Node | None = None): + def add( + self, node: torch.fx.Node, desc: str, dep: torch.fx.Node | None = None + ) -> None: """ Add a new event to the tracker. """ @@ -197,7 +198,13 @@ def add(self, node: torch.fx.Node, desc: str, dep: torch.fx.Node | None = None): self.node_events[node.name] = [] self.node_events[node.name].append(len(self.events) - 1) - def print_node(self, node_name, recursive=False, tab="", writer=None): + def print_node( + self, + node_name: str, + recursive: bool = False, + tab: str = "", + writer: Callable[[str], object] | None = None, + ) -> None: """ Print a node and its events. @param recursive: if True, print nodes that caused the events on this current node. @@ -214,7 +221,7 @@ def print_node(self, node_name, recursive=False, tab="", writer=None): event.dep.name, recursive=True, tab="| " + tab, writer=writer ) - def to_dict(self): + def to_dict(self) -> dict[str, list[str]]: """ Create dict dump on all events. """ @@ -226,7 +233,7 @@ def to_dict(self): ret[name].append(event.to_str()) return ret - def print_all(self, writer=None): + def print_all(self, writer: Callable[[str], object] | None = None) -> None: """ Print all nodes in a list. @param writer: function to write to file. If None, use print. @@ -237,7 +244,7 @@ def print_all(self, writer=None): writer(f"Node: {name}:") self.print_node(name, recursive=False, tab=" ", writer=writer) - def dump(self): + def dump(self) -> None: """ Function to be invoked at the end of the finder execution to printout tracked events specified by the mode. """ @@ -251,8 +258,8 @@ def dump(self): payload_fn=lambda: json.dumps(self.to_dict()), ) - def writeln(f): - def fn(x): + def make_writer(f: IO[str]) -> Callable[[str], int]: + def fn(x: str) -> int: return f.write(x + "\n") return fn @@ -261,16 +268,15 @@ def fn(x): # Mode >=1: Dump all events to file if self.tracker_mode >= 1: with open(self.dump_prefix + ALL_SUFFIX, "w") as f: - self.print_all(writeln(f)) + self.print_all(make_writer(f)) - def dump_selected_nodes(nodes): + def dump_selected_nodes(nodes: list[str]) -> None: with open(self.dump_prefix + NODES_SUFFIX, "w") as f: + writer = make_writer(f) for node_name in nodes: - writeln(f"===== Tracking node {node_name} =====") - self.print_node( - node_name, recursive=True, tab="|-", writer=writeln(f) - ) - writeln(f"===== End of tracking node {node_name} =====") + writer(f"===== Tracking node {node_name} =====") + self.print_node(node_name, recursive=True, tab="|-", writer=writer) + writer(f"===== End of tracking node {node_name} =====") # Mode 2: Dump specific nodes in recursive manner. # Mode 3: Dump all nodes with more than 1 event in recursive manner. @@ -307,7 +313,7 @@ def __init__( module: torch.fx.GraphModule, operator_support: OperatorSupportBase, allow_non_tensor: bool, - ): + ) -> None: self.module = module self.operator_support = operator_support self.allow_non_tensor = allow_non_tensor @@ -315,7 +321,7 @@ def __init__( self.tracker = NodeEventTracker(int(TRACKER_MODE), DUMP_PREFIX) - def reduce_acc_nodes_non_tensor_input_helper(self, cpu_worklist: NodeList): + def reduce_acc_nodes_non_tensor_input_helper(self, cpu_worklist: NodeList) -> None: """ Transitively excludes nodes from ACC supported set. For every node in the worklist: @@ -334,7 +340,7 @@ def reduce_acc_nodes_non_tensor_input_helper(self, cpu_worklist: NodeList): self.tracker.add(user, "new_cpu_node|non_tensor_output") cpu_worklist.append(user) - def reduce_acc_nodes_non_tensor_input(self): + def reduce_acc_nodes_non_tensor_input(self) -> None: """ Excludes nodes from ACC supported set that have direct upstream CPU nodes that produce non-tensor outputs. @@ -353,7 +359,7 @@ def reduce_acc_nodes_non_tensor_input(self): self.reduce_acc_nodes_non_tensor_input_helper(non_tensor_cpu_nodes) - def reduce_acc_nodes_non_tensor_output(self): + def reduce_acc_nodes_non_tensor_output(self) -> None: """ Excludes nodes from ACC supported set that produce non-tensor outputs and have downstream CPU nodes. @@ -452,10 +458,10 @@ def generate_inputs_for_submodules( """ handles = [] - results = {} + results: dict[str, Any] = {} submodule_to_names = {mod: name for name, mod in model.named_modules()} - def pre_forward(module, module_inputs): + def pre_forward(module: torch.nn.Module, module_inputs: tuple[Any, ...]) -> None: results[submodule_to_names[module]] = ( copy.deepcopy(module_inputs) if deepcopy else module_inputs ) @@ -465,7 +471,7 @@ def pre_forward(module, module_inputs): if not isinstance(mod, torch.jit.ScriptModule): handles.append(mod.register_forward_pre_hook(pre_forward)) - def clean_up_handles(): + def clean_up_handles() -> None: for h in handles: h.remove() @@ -535,7 +541,7 @@ def __init__( non_acc_submodule_name: str = "_run_on_cpu_", return_tuple: bool = False, nodes_finder: FxNetAccNodesFinder | None = None, - ): + ) -> None: """ Preprocesses graph before splitting: - finds nodes supported by ACC, @@ -606,7 +612,7 @@ def find_deps(self) -> dict[torch.fx.Node, NodeSet]: deps[user].add(node) return deps - def update_deps_for_fusions(self): + def update_deps_for_fusions(self) -> None: """ Updates graph of dependencies so that: - nodes from the same fusion depend on the same set of outer nodes, @@ -621,7 +627,7 @@ def update_deps_for_fusions(self): if user not in fusion: self.deps[user].add(node) - def _merge_overlapping_fusions(self): + def _merge_overlapping_fusions(self) -> None: """ Merge fusion groups that share nodes. @@ -709,7 +715,7 @@ def _find_culprit(self, mod: torch.fx.GraphModule, inputs: Tensors) -> str: def _draw_graph_based_on_node_support( self, mod: torch.fx.GraphModule, supported_nodes: NodeList - ): + ) -> None: color_map = { "default": "AliceBlue", "supported": "chartreuse1", @@ -717,7 +723,7 @@ def _draw_graph_based_on_node_support( } class CustomDrawer(FxGraphDrawer): - def _get_node_style(self, node): + def _get_node_style(self, node: torch.fx.Node) -> dict[str, str]: template = super()._get_node_style(node) if node in supported_nodes: template["fillcolor"] = color_map["supported"] @@ -733,14 +739,14 @@ def _get_node_style(self, node): # pyre-fixme[16]: `pydot.Dot` has no attribute `write_raw`. dot_graph.write_raw("node_support.dot") # type: ignore[attr-defined] - def node_support_preview(self, dump_graph: bool = False): + def node_support_preview(self, dump_graph: bool = False) -> str: submodules = dict(self.module.named_modules()) supported_nodes: NodeList = [] supported_node_types = defaultdict(set) unsupported_node_types = defaultdict(set) - def get_dtype(arg): + def get_dtype(arg: torch.fx.Node) -> torch.dtype | None: tensor_meta = arg.meta.get("tensor_meta") return getattr(tensor_meta, "dtype", None) @@ -800,7 +806,7 @@ def get_dtype(arg): # Return reports for testing purpose return reports - def split_preview(self, dump_graph: bool = False): + def split_preview(self, dump_graph: bool = False) -> str: reports = "" subgraphs = self.put_nodes_into_subgraphs() acc_subgraphs_num = len([g for g in subgraphs if g.is_acc]) @@ -842,16 +848,24 @@ def split_preview(self, dump_graph: bool = False): submod = getattr(split_mod, node.target) - def get_submod_inputs(main_mod, submod, example_inputs): - sub_inputs = None + def get_submod_inputs( + main_mod: torch.fx.GraphModule, + submod: torch.fx.GraphModule, + example_inputs: Sequence[Any], + ) -> tuple[Any, ...]: + sub_inputs: tuple[Any, ...] | None = None - def get_inputs(self, inputs): + def get_inputs( + self: torch.nn.Module, inputs: tuple[Any, ...] + ) -> None: nonlocal sub_inputs sub_inputs = inputs handle = submod.register_forward_pre_hook(get_inputs) main_mod(*example_inputs) handle.remove() + if sub_inputs is None: + raise AssertionError("Forward pre-hook did not capture inputs") return sub_inputs submod_inputs = get_submod_inputs(split_mod, submod, self.sample_input) @@ -872,7 +886,7 @@ def get_inputs(self, inputs): reports += "Checking outputs...\n" - def get_bytes(node: torch.fx.Node): + def get_bytes(node: torch.fx.Node) -> None: nonlocal total_output_bytes nonlocal reports if not is_node_output_tensor(node): @@ -937,7 +951,9 @@ def find_reverse_deps( return result - def update_reverse_deps_for_fusions(self, deps: dict[torch.fx.Node, NodeSet]): + def update_reverse_deps_for_fusions( + self, deps: dict[torch.fx.Node, NodeSet] + ) -> None: processed_node = set() for node, fusion in self.fusions.items(): @@ -981,7 +997,7 @@ def find_parent_nodes_of_subgraph(self, tag: str) -> NodeSet: return parent_nodes - def extend_acc_subgraph(self, tag: str): + def extend_acc_subgraph(self, tag: str) -> None: """ Extend the acc subgraph with `tag` going the reversed topological direction. """ @@ -1145,7 +1161,7 @@ def remove_small_acc_subgraphs(self, subgraphs: list[Subgraph]) -> list[Subgraph result.append(subgraph) return result - def tag(self, subgraphs: list[Subgraph]): + def tag(self, subgraphs: list[Subgraph]) -> None: self.tags = [] for subgraph in subgraphs: tag = ( @@ -1186,7 +1202,7 @@ def __call__(self) -> torch.fx.GraphModule: def generate_split_results(self) -> SplitResult: split_module = self() - submodule_names = [] + submodule_names: list[str] = [] for name, _mod in split_module.named_children(): submodule_names.append(name) if ( diff --git a/torch/fx/passes/tests/_test_split_utils.py b/torch/fx/passes/tests/_test_split_utils.py index d841685668624..e1751d7f88686 100644 --- a/torch/fx/passes/tests/_test_split_utils.py +++ b/torch/fx/passes/tests/_test_split_utils.py @@ -42,10 +42,12 @@ def _create_mock_node( node.meta = {"type": int} # Non-tensor type # Mock users dict (Node.users is dict[Node, None]) - node.users = {} + users: dict[torch.fx.Node, None] = {} + node.users = users # Initialize the _input_nodes dict (Node._input_nodes is dict[Node, None]) - node._input_nodes = {} + input_nodes: dict[torch.fx.Node, None] = {} + node._input_nodes = input_nodes return node diff --git a/torch/fx/passes/tools_common.py b/torch/fx/passes/tools_common.py index 75f76a6d0da18..751a615698cbd 100644 --- a/torch/fx/passes/tools_common.py +++ b/torch/fx/passes/tools_common.py @@ -1,4 +1,3 @@ -# mypy: allow-untyped-defs import collections import heapq import operator @@ -30,7 +29,7 @@ @compatibility(is_backward_compatible=False) -def get_acc_ops_name(k): +def get_acc_ops_name(k: str | type) -> str: if isinstance(k, str): return k elif k.__module__ and "acc_ops" in k.__module__: @@ -105,7 +104,7 @@ class FxNetAccFusionsFinder: Such groups are called fusion groups. """ - def __init__(self, module: torch.fx.GraphModule, acc_nodes: NodeSet): + def __init__(self, module: torch.fx.GraphModule, acc_nodes: NodeSet) -> None: self.module = module self.nodes = list(module.graph.nodes) self.acc_nodes = acc_nodes @@ -125,7 +124,7 @@ class FusionGroup: # Nodes that in the fusion group that haven't been processed yet. nodes_need_process: NodeSet - def add_node(self, node): + def add_node(self, node: torch.fx.Node) -> None: """ Add a node to fusion group. """ @@ -148,7 +147,7 @@ def recursive_add_node( fusion_group: "FxNetAccFusionsFinder.FusionGroup", inputs: NodeSet | NodeList, visited: NodeSet | None = None, - ): + ) -> bool: """ Start from inputs and going reverse topological order. If any upstream node is in the fusion group, add all the nodes in this path to fusion group. @@ -302,7 +301,7 @@ def legalize_graph(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: for node in gm.graph.nodes: for user in node.users: indeg[user] += 1 - queue: collections.deque = collections.deque() + queue: collections.deque[torch.fx.Node] = collections.deque() # Add all nodes with no dependencies to the queue for node in gm.graph.nodes: if indeg[node] == 0: diff --git a/torch/fx/passes/utils/common.py b/torch/fx/passes/utils/common.py index a0a375b96dac2..87f9fa4326845 100644 --- a/torch/fx/passes/utils/common.py +++ b/torch/fx/passes/utils/common.py @@ -1,5 +1,3 @@ -# mypy: allow-untyped-defs - from torch.fx._compatibility import compatibility from torch.fx.graph import Graph from torch.fx.graph_module import GraphModule @@ -17,7 +15,7 @@ class HolderModule(Module): that uses the attributes """ - def __init__(self, d): + def __init__(self, d: dict[str, Module | None]) -> None: super().__init__() for k, v in d.items(): self.add_module(k, v) diff --git a/torch/fx/passes/utils/matcher_utils.py b/torch/fx/passes/utils/matcher_utils.py index e0ebe687e95f5..c5c59e40e3007 100644 --- a/torch/fx/passes/utils/matcher_utils.py +++ b/torch/fx/passes/utils/matcher_utils.py @@ -1,4 +1,3 @@ -# mypy: allow-untyped-defs import copy import logging import os @@ -15,7 +14,7 @@ # Set`PYTORCH_MATCHER_LOGLEVEL=INFO` to see debug logs -def _init_logger(): +def _init_logger() -> logging.Logger: logger = logging.getLogger(__name__) level = os.environ.get("PYTORCH_MATCHER_LOGLEVEL", "WARNING").upper() @@ -51,7 +50,7 @@ class InternalMatch: # only available if the matcher is `SubgraphMatcherWithNameNodesMap` name_node_map: dict[str, Node] = field(default_factory=dict) - def __copy__(self): + def __copy__(self) -> "InternalMatch": return InternalMatch( anchors=self.anchors, nodes_map=self.nodes_map.copy(), @@ -251,7 +250,9 @@ def _match_nodes( # match for `gn` match_found = True - def _match_args(args1: list | tuple, args2: list | tuple) -> bool: + def _match_args( + args1: list[Any] | tuple[Any, ...], args2: list[Any] | tuple[Any, ...] + ) -> bool: if len(args1) != len(args2): return False @@ -271,7 +272,8 @@ def _match_args(args1: list | tuple, args2: list | tuple) -> bool: return True # Flatten all args/kwargs into 1 list of args - pn_args, gn_args = None, None + pn_args: list[Any] | None = None + gn_args: list[Any] | None = None if ( ( len(pn.args) != len(gn.args) @@ -282,7 +284,9 @@ def _match_args(args1: list | tuple, args2: list | tuple) -> bool: ): args_schema = pn.target._schema.arguments - def get_all_arguments(orig_args, orig_kwargs): + def get_all_arguments( + orig_args: tuple[Any, ...], orig_kwargs: dict[str, Any] + ) -> list[Any]: all_args = [] for i, schema in enumerate(args_schema): if schema.name in orig_kwargs: @@ -370,7 +374,7 @@ def match(self, graph: Graph, node_name_match: str = "") -> list[InternalMatch]: matches: list[InternalMatch] = [] - def backtracking(anchor_index, match): + def backtracking(anchor_index: int, match: InternalMatch) -> None: if anchor_index == len(match_candidates_list): match.placeholder_nodes = [ match.nodes_map[pn] for pn in self.pattern_placeholder_nodes @@ -418,7 +422,7 @@ def backtracking(anchor_index, match): ) # filter out the matches that form a cycle if the subgraph is fused - valid_matches = [] + valid_matches: list[InternalMatch] = [] for match in matches: matched_compute_nodes = [ gn diff --git a/torch/fx/passes/utils/matcher_with_name_node_map_utils.py b/torch/fx/passes/utils/matcher_with_name_node_map_utils.py index 8b2f0f2e89e22..43726f79d1ed7 100644 --- a/torch/fx/passes/utils/matcher_with_name_node_map_utils.py +++ b/torch/fx/passes/utils/matcher_with_name_node_map_utils.py @@ -13,7 +13,7 @@ def _split_to_graph_and_name_node_map( from torch.fx.graph import _PyTreeInfo from torch.utils._pytree import tree_flatten, tree_unflatten - name_node_map = {} + name_node_map: dict[str, Node] = {} for n in gm.graph.nodes: if n.op == "output": if gm._out_spec is None: diff --git a/torch/fx/proxy.py b/torch/fx/proxy.py index 0d0e13707c9d1..f1e86f55bfba5 100644 --- a/torch/fx/proxy.py +++ b/torch/fx/proxy.py @@ -1,5 +1,3 @@ -# mypy: ignore-errors - import collections import copy import dis @@ -9,10 +7,12 @@ import operator import sys import traceback +import types from collections import OrderedDict from collections.abc import Callable, Iterator from dataclasses import fields, is_dataclass -from typing import Any, TypeVar +from typing import Any, cast, TypeVar +from typing_extensions import Never import torch import torch.fx.traceback as fx_traceback @@ -117,7 +117,7 @@ def forward(self, x): """ - def __init__(self, module_path: str, module_type: Any): + def __init__(self, module_path: str, module_type: Any) -> None: super().__init__() self.module_path = module_path self.module_type = module_type @@ -134,7 +134,7 @@ def __init__( self, scope: Scope, current_scope: Scope, - ): + ) -> None: super().__init__() # Keep a copy of prev scope to restore on exit self._prev_scope = copy.copy(scope) @@ -144,10 +144,10 @@ def __init__( # Save a reference so we can restore it self._scope = scope - def __enter__(self): + def __enter__(self) -> Scope: return self._scope - def __exit__(self, *args): + def __exit__(self, *args: object) -> None: self._scope.module_path = self._prev_scope.module_path self._scope.module_type = self._prev_scope.module_type return @@ -217,7 +217,8 @@ def create_node( """ if kind == "call_function" and self.check_mutable_operations: - check_for_mutable_operation(target, args, kwargs) + target_fn = cast(Callable[..., Any], target) + check_for_mutable_operation(target_fn, args, kwargs) node = self.graph.create_node(kind, target, args, kwargs, name, type_expr) # TODO node_name_to_scope will be depreciated in favor of @@ -254,7 +255,7 @@ def create_node( # See Note [Functionalization View Replay Annotation] # Overriding some node meta with the original node meta of the # regenerated node. - replay_node: Node = fx_traceback.get_current_replay_node() + replay_node: Node | None = fx_traceback.get_current_replay_node() if replay_node is not None: node.meta["is_functional_regenerated"] = True if "custom" in replay_node.meta: @@ -280,7 +281,7 @@ def _filter_traceback_frames( ) -> traceback.StackSummary: # This method can be overridden to customize the frame filtering logic # for the recorded stack trace - user_frames = [] + user_frames: list[traceback.FrameSummary] = [] if self._record_forward_stack_traces_only: user_frames = [ frame @@ -301,7 +302,7 @@ def _filter_traceback_frames( # Not having a "forward" call in the stacktrace implies the # stacktrace will probably be irrelevant if first_forward == -1: - user_frames = [] + user_frames: list[traceback.FrameSummary] = [] from torch.fx.experimental.symbolic_shapes import uninteresting_files @@ -326,9 +327,8 @@ def create_proxy( kwargs: dict[str, Any], name: str | None = None, type_expr: Any | None = None, - # fix noqa when updating bc tests - proxy_factory_fn: Callable[[Node], "Proxy"] = None, # noqa: RUF013 - ): + proxy_factory_fn: Callable[[Node], "Proxy"] | None = None, + ) -> "Proxy": """ Create a Node from the given arguments, then return the Node wrapped in a Proxy object. @@ -355,7 +355,7 @@ def create_proxy( return proxy - def _find_user_frame(self): + def _find_user_frame(self) -> types.FrameType | None: """ Find the Python stack frame executing the user code during symbolic tracing. @@ -440,14 +440,14 @@ def create_arg(self, a: Any) -> Argument: ) elif isinstance(a, range): - return range( - self.create_arg(a.start), - self.create_arg(a.stop), - self.create_arg(a.step), + return range( # pyrefly: ignore[no-matching-overload] + self.create_arg(a.start), # pyrefly: ignore[bad-argument-type] + self.create_arg(a.stop), # pyrefly: ignore[bad-argument-type] + self.create_arg(a.step), # pyrefly: ignore[bad-argument-type] ) elif isinstance(a, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)): - return a + return a # pyrefly: ignore[bad-return] elif is_opaque_value_type(type(a)): return a @@ -460,7 +460,7 @@ def create_arg(self, a: Any) -> Argument: return self.create_node("call_function", a.__class__, (), kwargs) elif isinstance(a, (*base_types, enum.Enum)) or a is None or a is ...: - return a + return a # pyrefly: ignore[bad-return] raise NotImplementedError(f"argument of type: {type(a)}") @@ -494,7 +494,7 @@ def iter(self, obj: "Proxy") -> Iterator: ) @compatibility(is_backward_compatible=True) - def keys(self, obj: "Proxy") -> Any: + def keys(self, obj: "Proxy") -> "Proxy": """Called when a proxy object is has the keys() method called. This is what happens when ** is called on a proxy. This should return an iterator it ** is suppose to work in your custom tracer. @@ -502,7 +502,7 @@ def keys(self, obj: "Proxy") -> Any: return Attribute(obj, "keys")() -def _get_seq_nr(node_name: str = ""): +def _get_seq_nr(node_name: str = "") -> int | None: """ Returns the seq_nr node meta for the current proxy node that we're creating. The seq_nr number in node meta is related to but not the same as the "sequence number" @@ -547,7 +547,7 @@ def _get_seq_nr(node_name: str = ""): # See Note [Functionalization View Replay Annotation] # Overriding some node meta with the original node meta of the # regenerated node. - replay_node: Node = fx_traceback.get_current_replay_node() + replay_node: Node | None = fx_traceback.get_current_replay_node() if replay_node is not None: if "seq_nr" in replay_node.meta: annotation_log.debug("%s: seq_nr from replay_node", node_name) @@ -559,16 +559,16 @@ def _get_seq_nr(node_name: str = ""): # used in Proxy object when just appending to the graph while not tracing. @compatibility(is_backward_compatible=True) class GraphAppendingTracer(TracerBase): - def __init__(self, graph: Graph): + def __init__(self, graph: Graph) -> None: super().__init__() self.graph = graph self.scope = Scope("", None) - self.module_stack = collections.OrderedDict() - self.node_name_to_scope = {} + self.module_stack: OrderedDict[str, tuple[str, Any]] = collections.OrderedDict() + self.node_name_to_scope: dict[str, tuple[str, type]] = {} @compatibility(is_backward_compatible=False) -def assert_fn(x): +def assert_fn(x: object) -> None: if not x: raise AssertionError("Assertion failed") @@ -609,7 +609,7 @@ class Proxy: """ @compatibility(is_backward_compatible=True) - def __init__(self, node: Node, tracer: "TracerBase | None" = None): + def __init__(self, node: Node, tracer: "TracerBase | None" = None) -> None: if tracer is None: # This allows you to create a Proxy object around a raw Node tracer = GraphAppendingTracer(node.graph) @@ -619,21 +619,21 @@ def __init__(self, node: Node, tracer: "TracerBase | None" = None): def __repr__(self) -> str: return f"Proxy({self.node.name})" - def __getattr__(self, k) -> "Attribute": + def __getattr__(self, k: str) -> "Attribute": # note: not added to the graph yet, if this is a method call # we peephole optimize to the method invocation return Attribute(self, k) - def __getstate__(self) -> dict: + def __getstate__(self) -> dict[str, Any]: return self.__dict__ - def __deepcopy__(self, memo) -> dict: + def __deepcopy__(self, memo: dict[int, Any]) -> "Proxy": # We have to explicitly override this method, because otherwise deepcopy # will go to __getattr__(self, "__deepcopy__") and return a # Attribute(__deepcopy__), and may go into an infinite loop in some cases. import copy - new_dict = {} + new_dict: dict[str, Any] = {} for k, v in self.__dict__.items(): try: new_obj = copy.deepcopy(v, memo) @@ -655,11 +655,11 @@ def __deepcopy__(self, memo) -> dict: new_proxy.__dict__[k] = v return new_proxy - def __setstate__(self, d): + def __setstate__(self, d: dict[str, Any]) -> None: # This is called when being unpickled/loaded. self.__dict__ = d - def __call__(self, *args, **kwargs) -> "Proxy": + def __call__(self, *args: Any, **kwargs: Any) -> "Proxy": return self.tracer.create_proxy( "call_method", "__call__", (self,) + args, kwargs ) @@ -686,7 +686,7 @@ def __iter__(self) -> Iterator["Proxy"]: return self.tracer.iter(self) - def __abs__(self): + def __abs__(self) -> "Proxy": return self.tracer.create_proxy("call_function", operator.abs, (self,), {}) def __bool__(self) -> bool: @@ -725,10 +725,10 @@ def __bool__(self) -> bool: return self.tracer.to_bool(self) @compatibility(is_backward_compatible=True) - def keys(self): + def keys(self) -> "Proxy": return self.tracer.keys(self) - def __len__(self): + def __len__(self) -> int: raise RuntimeError( "'len' is not supported in symbolic tracing by default. If you want " "this call to be recorded, please call torch.fx.wrap('len') at " @@ -736,13 +736,19 @@ def __len__(self): ) @classmethod - def __torch_function__(cls, orig_method, types, args=None, kwargs=None): + def __torch_function__( + cls, + orig_method: Callable[..., Any], + types: tuple[type, ...], + args: tuple[Any, ...] | None = None, + kwargs: dict[str, Any] | None = None, + ) -> "Proxy": args = args if args else () kwargs = kwargs if kwargs else {} - tracers: dict[Any, None] = {} + tracers: dict[TracerBase, None] = {} - def find_tracer(a): + def find_tracer(a: Any) -> None: if isinstance(a, cls): tracers[a.tracer] = None @@ -788,7 +794,9 @@ class MetaProxy(Proxy): A Proxy subclass that propagates metadata (meta['val']) during graph tracing. """ - def __init__(self, node: Node, tracer: "TracerBase | None" = None, fake_mode=None): + def __init__( + self, node: Node, tracer: "TracerBase | None" = None, fake_mode: Any = None + ) -> None: super().__init__(node, tracer) self.fake_mode = fake_mode @@ -796,7 +804,13 @@ def __repr__(self) -> str: return f"MetaProxy({self.node.name})" @classmethod - def __torch_function__(cls, orig_method, types, args=None, kwargs=None): + def __torch_function__( + cls, + orig_method: Callable[..., Any], + types: tuple[type, ...], + args: tuple[Any, ...] | None = None, + kwargs: dict[str, Any] | None = None, + ) -> "MetaProxy": args = args if args else () kwargs = kwargs if kwargs else {} @@ -823,14 +837,14 @@ def __torch_function__(cls, orig_method, types, args=None, kwargs=None): @compatibility(is_backward_compatible=True) class Attribute(Proxy): @compatibility(is_backward_compatible=True) - def __init__(self, root: Proxy, attr: str): + def __init__(self, root: Proxy, attr: str) -> None: self.root = root self.attr = attr self.tracer = root.tracer self._node: Node | None = None @property - def node(self): + def node(self) -> Node: # pyrefly: ignore[bad-override] # the node for attributes is added lazily, since most will just be method calls # which do not rely on the getitem call if self._node is None: @@ -839,7 +853,7 @@ def node(self): ).node return self._node - def __call__(self, *args, **kwargs): + def __call__(self, *args: Any, **kwargs: Any) -> "Proxy": return self.tracer.create_proxy( "call_method", self.attr, (self.root,) + args, kwargs ) @@ -853,7 +867,9 @@ class ParameterProxy(Proxy): so that conditional tests on these attributes will not throw exception during tracing """ - def __init__(self, tracer: TracerBase, node: Node, name, param): + def __init__( + self, tracer: TracerBase, node: Node, name: str, param: torch.nn.Parameter + ) -> None: super().__init__(node, tracer) if not isinstance(param, torch.nn.Parameter): raise AssertionError(f"Expected Parameter, got {type(param)}") @@ -864,30 +880,30 @@ def __repr__(self) -> str: return f"ParameterProxy({self.name})" @property - def shape(self): + def shape(self) -> torch.Size: return self.param.shape - def size(self): + def size(self) -> torch.Size: return self.param.size() - def dim(self): + def dim(self) -> int: return self.param.dim() @property - def ndim(self): + def ndim(self) -> int: return self.param.ndim - def numel(self): + def numel(self) -> int: return self.param.numel() - def nelement(self): + def nelement(self) -> int: return self.param.nelement() for method in magic_methods: - def _scope(method): - def impl(*args, **kwargs): + def _scope(method: str) -> None: + def impl(*args: Any, **kwargs: Any) -> "Proxy": tracer = args[0].tracer target = getattr(operator, method) return tracer.create_proxy("call_function", target, args, kwargs) @@ -899,10 +915,10 @@ def impl(*args, **kwargs): _scope(method) -def _define_reflectable(orig_method_name): +def _define_reflectable(orig_method_name: str) -> None: method_name = f"__r{orig_method_name.strip('_')}__" - def impl(self, rhs): + def impl(self: "Proxy", rhs: Any) -> "Proxy": target = getattr(operator, orig_method_name) return self.tracer.create_proxy("call_function", target, (rhs, self), {}) @@ -915,15 +931,15 @@ def impl(self, rhs): _define_reflectable(orig_method_name) -def _no_nodes_error(arg): +def _no_nodes_error(arg: Argument) -> Never: raise RuntimeError( "Keys for dictionaries used as an argument cannot contain a " f"Node. Got key: {arg}" ) -def _create_arg_dict(self, a): - r = {} +def _create_arg_dict(self: TracerBase, a: dict[Any, Any]) -> dict[Any, Argument]: + r: dict[Any, Argument] = {} for k, v in a.items(): if not isinstance(k, str): # Check for invalid dict keys. We do not want a Proxy to appear diff --git a/torch/fx/subgraph_rewriter.py b/torch/fx/subgraph_rewriter.py index 564c0aeadf8d6..71d884cc86e70 100644 --- a/torch/fx/subgraph_rewriter.py +++ b/torch/fx/subgraph_rewriter.py @@ -95,8 +95,8 @@ def try_get_attr(gm: torch.nn.Module, target: str) -> Any | None: @compatibility(is_backward_compatible=True) def replace_pattern( gm: GraphModule, - pattern: Callable | GraphModule, - replacement: Callable | GraphModule, + pattern: Callable[..., Any] | GraphModule, + replacement: Callable[..., Any] | GraphModule, ) -> list[Match]: """ Matches all possible non-overlapping sets of operators and their @@ -225,8 +225,8 @@ def forward(self, x, w1, w2): @compatibility(is_backward_compatible=False) def replace_pattern_with_filters( gm: GraphModule, - pattern: Callable | Graph | GraphModule, - replacement: Callable | Graph | GraphModule | None = None, + pattern: Callable[..., Any] | Graph | GraphModule, + replacement: Callable[..., Any] | Graph | GraphModule | None = None, match_filters: list[Callable[["InternalMatch", Graph, Graph], bool]] | None = None, ignore_literals: bool = False, # Placed at the end to avoid breaking backward compatibility @@ -261,8 +261,8 @@ def replace_pattern_with_filters( def _replace_pattern( gm: GraphModule, - pattern: Callable | Graph | GraphModule, - replacement: Callable | Graph | GraphModule | None = None, + pattern: Callable[..., Any] | Graph | GraphModule, + replacement: Callable[..., Any] | Graph | GraphModule | None = None, match_filters: list[Callable[["InternalMatch", Graph, Graph], bool]] | None = None, ignore_literals: bool = False, # Placed at the end to avoid breaking backward compatibility @@ -417,7 +417,9 @@ def _replace_pattern( f"{len(copied_returning_nodes)}" # pyrefly: ignore [bad-argument-type] ) for gn, copied_node in zip(match.returning_nodes, copied_returning_nodes): # type: ignore[arg-type] + # pyrefly: ignore [bad-argument-type] gn.replace_all_uses_with(copied_node) + # pyrefly: ignore [unsupported-operation] match_changed_node[gn] = copied_node # Remove the original nodes for node in reversed(pattern_graph.nodes): diff --git a/torch/fx/tensor_type.py b/torch/fx/tensor_type.py index 4f375e461ef28..f8342146e4b8a 100644 --- a/torch/fx/tensor_type.py +++ b/torch/fx/tensor_type.py @@ -1,9 +1,13 @@ -# mypy: allow-untyped-defs +from collections.abc import Sequence + from torch.fx.experimental.unification import Var # type: ignore[attr-defined] from ._compatibility import compatibility +__all__ = ["Dyn", "TensorType", "is_consistent", "is_more_precise"] + + @compatibility(is_backward_compatible=False) class TensorType: """ @@ -14,21 +18,21 @@ def forward(self, x:TensorType((1,2,3, Dyn)), y:TensorType((1,2,3, Dyn))): return torch.add(x, y) """ - def __init__(self, dim): + def __init__(self, dim: Sequence[object]) -> None: self.__origin__ = TensorType self.__args__ = dim - def __repr__(self): + def __repr__(self) -> str: return f"TensorType[{self.__args__}]" - def __eq__(self, other): + def __eq__(self, other: object) -> bool: if isinstance(other, self.__class__): return list(self.__args__) == list(other.__args__) else: return False @staticmethod - def __class_getitem__(*args): + def __class_getitem__(*args: object) -> "TensorType": if len(args) == 1 and isinstance(args[0], tuple): args = args[0] return TensorType(tuple(args)) @@ -42,13 +46,13 @@ class _DynType: def __init__(self) -> None: self.__name__ = "_DynType" - def __eq__(self, other): + def __eq__(self, other: object) -> bool: return isinstance(other, self.__class__) - def __str__(self): + def __str__(self) -> str: return "Dyn" - def __repr__(self): + def __repr__(self) -> str: return "Dyn" @@ -56,7 +60,7 @@ def __repr__(self): @compatibility(is_backward_compatible=False) -def is_consistent(t1, t2): +def is_consistent(t1: object, t2: object) -> bool: """ A binary relation denoted by ~ that determines if t1 is consistent with t2. The relation is reflexive, symmetric but not transitive. @@ -84,7 +88,7 @@ def is_consistent(t1, t2): @compatibility(is_backward_compatible=False) -def is_more_precise(t1, t2): +def is_more_precise(t1: object, t2: object) -> bool: """ A binary relation denoted by <= that determines if t1 is more precise than t2. The relation is reflexive and transitive. diff --git a/torch/fx/traceback.py b/torch/fx/traceback.py index f13c2837b4344..b6a9ddee2dd76 100644 --- a/torch/fx/traceback.py +++ b/torch/fx/traceback.py @@ -1,11 +1,11 @@ -# mypy: allow-untyped-defs import copy import logging import traceback from collections import defaultdict +from collections.abc import Callable, Iterator from contextlib import contextmanager from enum import Enum -from typing import Any, Optional, Union +from typing import Any, Optional, ParamSpec, TypeVar, Union from torch._utils_internal import signpost_event @@ -15,6 +15,9 @@ from .node import Node +_P = ParamSpec("_P") +_R = TypeVar("_R") + log = logging.getLogger(__name__) __all__ = [ @@ -82,7 +85,7 @@ class NodeSource: """ class NodeInfo: - def __init__(self, name: str, target: str, graph_id: int): + def __init__(self, name: str, target: str, graph_id: int) -> None: self.name = name self.target = target self.graph_id = graph_id @@ -99,7 +102,7 @@ def __init__( node: Node | None, pass_name: str = "", action: Union["NodeSourceAction", list["NodeSourceAction"]] | None = None, - ): + ) -> None: self.pass_name = pass_name if action is None: @@ -139,15 +142,15 @@ def target(self) -> str: def graph_id(self) -> int: return self.node_info.graph_id if self.node_info else -1 - def __repr__(self): + def __repr__(self) -> str: return self.print_readable() - def _get_action_string(self): + def _get_action_string(self) -> str: if self._action_string is None: self._action_string = "+".join([a.name.lower() for a in self.action]) return self._action_string - def print_readable(self, indent=0): + def print_readable(self, indent: int = 0) -> str: if indent > 9: return "" result = "" @@ -160,7 +163,7 @@ def print_readable(self, indent=0): result += item.print_readable(indent + 1) return result - def to_dict(self) -> dict: + def to_dict(self) -> dict[str, Any]: if self._dict is None: # Convert the object to a dictionary action_string = self._get_action_string() @@ -177,15 +180,15 @@ def to_dict(self) -> dict: raise AssertionError("_dict is None after initialization") return self._dict - def __eq__(self, other: object): + def __eq__(self, other: object) -> bool: if not isinstance(other, NodeSource): return False return self.to_dict() == other.to_dict() - def __hash__(self): + def __hash__(self) -> int: # Create a hash based on the dictionary representation # We need to convert the dict to a hashable form - def _make_hashable(obj): + def _make_hashable(obj: Any) -> Any: if isinstance(obj, dict): return tuple(sorted((k, _make_hashable(v)) for k, v in obj.items())) elif isinstance(obj, list): @@ -196,7 +199,7 @@ def _make_hashable(obj): return hash(_make_hashable(self.to_dict())) @classmethod - def _from_dict(cls, d: dict | None) -> Optional["NodeSource"]: + def _from_dict(cls, d: dict[str, Any] | None) -> Optional["NodeSource"]: """ Recursively deserialize from_node metadata from dictionary data. It is used to deserialize the from_node field from serialized metadata. @@ -253,7 +256,7 @@ def _from_dict(cls, d: dict | None) -> Optional["NodeSource"]: @compatibility(is_backward_compatible=False) @contextmanager -def preserve_node_meta(enable=True): +def preserve_node_meta(enable: bool = True) -> Iterator[None]: global should_preserve_node_meta global current_meta saved_should_preserve_node_meta = should_preserve_node_meta @@ -268,7 +271,7 @@ def preserve_node_meta(enable=True): @contextmanager -def _preserve_node_seq_nr(preserve_seq_nr=True): +def _preserve_node_seq_nr(preserve_seq_nr: bool = True) -> Iterator[None]: """ Temporarily enables or disables the preservation of node.meta["seq_nr"] in the tracing context. @@ -284,7 +287,7 @@ def _preserve_node_seq_nr(preserve_seq_nr=True): @compatibility(is_backward_compatible=False) -def set_stack_trace(stack: list[str]): +def set_stack_trace(stack: list[str]) -> None: global current_meta if should_preserve_node_meta: @@ -298,7 +301,7 @@ def set_stack_trace(stack: list[str]): @compatibility(is_backward_compatible=False) @contextmanager -def annotate(annotation_dict: dict): +def annotate(annotation_dict: dict[str, Any]) -> Iterator[None]: """ Temporarily adds custom annotations to the current tracing context. The fx_node produced from this tracing context will have the @@ -338,7 +341,7 @@ def annotate(annotation_dict: dict): try: if not has_custom: - current_meta["custom"] = {} + current_meta["custom"] = dict[str, Any]() # Update with all key-value pairs from the input dict current_meta["custom"].update(annotation_dict) @@ -352,7 +355,9 @@ def annotate(annotation_dict: dict): @compatibility(is_backward_compatible=False) -def annotate_fn(annotation_dict: dict): +def annotate_fn( + annotation_dict: dict[str, Any], +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: """ A decorator that wraps a function with the annotate context manager. Use this when you want to annotate an entire function instead of a specific code block. @@ -377,19 +382,21 @@ def annotate_fn(annotation_dict: dict): """ from functools import wraps - def decorator(func): + def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]: @wraps(func) - def wrapper(*args, **kwargs): + # NB: Do not annotate with _P.args/_P.kwargs here. Dynamo guards on + # the identity of ParamSpec annotation objects, causing guard failures. + def wrapper(*args: Any, **kwargs: Any) -> Any: with annotate(annotation_dict): return func(*args, **kwargs) - return wrapper + return wrapper # type: ignore[return-value] return decorator @compatibility(is_backward_compatible=False) -def set_grad_fn_seq_nr(seq_nr): +def set_grad_fn_seq_nr(seq_nr: int) -> None: global current_meta if should_preserve_node_meta: @@ -401,7 +408,7 @@ def set_grad_fn_seq_nr(seq_nr): @compatibility(is_backward_compatible=False) -def reset_grad_fn_seq_nr(): +def reset_grad_fn_seq_nr() -> None: # NB: reset state properly, this would be helpful towards supporting # reentrant autograd if we actually wanted to do that. global current_meta @@ -437,7 +444,7 @@ def _is_preserving_node_seq_nr() -> bool: @compatibility(is_backward_compatible=False) @contextmanager -def set_current_meta(node, pass_name=""): +def set_current_meta(node: Node, pass_name: str = "") -> Iterator[None]: global current_meta if should_preserve_node_meta and node.meta: saved_meta = current_meta @@ -465,7 +472,7 @@ def get_current_meta() -> dict[str, Any]: @compatibility(is_backward_compatible=False) @contextmanager -def set_current_replay_node(node): +def set_current_replay_node(node: Node | None) -> Iterator[None]: """ Set the currently replay node. If `current_replay_node` is not None, then we're re-generating the `current_replay_node` in FunctionalTensorMode. @@ -481,7 +488,7 @@ def set_current_replay_node(node): @compatibility(is_backward_compatible=False) -def get_current_replay_node(): +def get_current_replay_node() -> Node | None: """ Get the currently replay node """ @@ -522,7 +529,7 @@ def _get_custom_metadata(gm: GraphModule) -> str: if not isinstance(gm, GraphModule): raise AssertionError(f"Expected GraphModule, got {type(gm)}") - def helper(gm: GraphModule): + def helper(gm: GraphModule) -> list[Any]: custom_metadata = [] for node in gm.graph.nodes: if hasattr(node, "meta") and node.meta.get("custom", None): @@ -530,7 +537,10 @@ def helper(gm: GraphModule): if node.op == "get_attr" and isinstance( getattr(gm, node.target), GraphModule ): - custom_metadata.append(helper(getattr(gm, node.target))) + custom_metadata.append( + # pyrefly: ignore[bad-argument-type] + helper(getattr(gm, node.target)) + ) return custom_metadata return "\n".join(str(x) for x in helper(gm)) diff --git a/torch/headeronly/CMakeLists.txt b/torch/headeronly/CMakeLists.txt index 93d2d7802b528..e6065f96f12e5 100644 --- a/torch/headeronly/CMakeLists.txt +++ b/torch/headeronly/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.27 FATAL_ERROR) project(headeronly CXX) -set(CMAKE_CXX_STANDARD 17 CACHE STRING "The C++ standard whose features are requested to build this target.") +set(CMAKE_CXX_STANDARD 20 CACHE STRING "The C++ standard whose features are requested to build this target.") set(CMAKE_EXPORT_COMPILE_COMMANDS ON) # Main build file for torch/headeronly, except there's no build cuz this lib is header-only! diff --git a/torch/headeronly/core/Dispatch_v2.h b/torch/headeronly/core/Dispatch_v2.h index 47ca18ce79c54..aa1bc1d9e9f72 100644 --- a/torch/headeronly/core/Dispatch_v2.h +++ b/torch/headeronly/core/Dispatch_v2.h @@ -54,6 +54,11 @@ torch::headeronly::ScalarType::UInt16, \ torch::headeronly::ScalarType::UInt32, \ torch::headeronly::ScalarType::UInt64 +#define AT_OPAQUE_TYPES \ + torch::headeronly::ScalarType::Byte, torch::headeronly::ScalarType::UInt16, \ + torch::headeronly::ScalarType::UInt32, \ + torch::headeronly::ScalarType::UInt64, \ + torch::headeronly::ScalarType::ComplexDouble #define AT_INTEGRAL_TYPES_V2 \ AT_EXPAND(AT_INTEGRAL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES) #define AT_COMPLEX_TYPES \ diff --git a/torch/headeronly/macros/Macros.h b/torch/headeronly/macros/Macros.h index 880e741abf62b..cef99df3f566f 100644 --- a/torch/headeronly/macros/Macros.h +++ b/torch/headeronly/macros/Macros.h @@ -325,41 +325,88 @@ constexpr uint32_t CUDA_THREADS_PER_BLOCK_FALLBACK = 256; #define C10_HIP_HOST_DEVICE #endif -#if defined(USE_ROCM) // C10_WARP_SIZE is only allowed for device code. -// Host code _must_ use at::cuda::warp_size() +// Host code dynamically-sized launch configs _must_ use at::cuda::warp_size(). +// Host or device statically-sized arrays _must_ use either +// C10_WARP_SIZE_UPPER_BOUND or C10_WARP_SIZE_LOWER_BOUND, as needed. +// // HIP header used to define warpSize as a constexpr that was either 32 or 64 // depending on the target device, and then always set it to 64 for host code. -// Host pass of HIP compiler needs C10_WARP_SIZE defined to _something_ so we -// set it to something unreasonable to trigger obvious host code errors. - +// For a time, that allowed C10_WARP_SIZE to be defined like so: +// +// #ifdef USE_ROCM +// #define C10_WARP_SIZE warpSize +// #else +// #define C10_WARP_SIZE 32 +// #endif +// +// In ROCm 7, warpSize is no longer constexpr, matching CUDA behavior. +// We can now only use warpSize for C10_WARP_SIZE in device code and this is +// enforced by using __device__ in its definition. In host code where +// C10_WARP_SIZE was previously used as a compile-time constant, this will now +// cause a compile-time error. +// +// If an array was previously expected to be sized at compile-time using +// C10_WARP_SIZE, users must now use either C10_WARP_SIZE_UPPER_BOUND or +// C10_WARP_SIZE_LOWER_BOUND depending on the situation. +// +// If C10_WARP_SIZE was previously used to determine kernel launch sizes, users +// must now use at::cuda::warp_size() for the dynamic runtime query. +// +// Unfortunately, C10_WARP_SIZE has been public and available for both host and +// device since approximately 2019, so forcing it to be device-only would break +// existing code in the wild. +#if defined(USE_ROCM) namespace at::cuda { TORCH_CUDA_CPP_API int warp_size(); } -#ifdef __HIPCC__ -static inline int __host__ C10_WARP_SIZE_INTERNAL() { +#if defined(__HIPCC__) +static __host__ inline int C10_WARP_SIZE_INTERNAL() { return at::cuda::warp_size(); } - -static inline constexpr int __device__ C10_WARP_SIZE_INTERNAL() { +// NOTE: __device__ C10_WARP_SIZE_INTERNAL +// For __SPIRV__, we must use dynamic warpSize. When not targeting __SPIRV__, +// we can use constexpr. This matches prior behavior. We preserve this for +// backward compatibility instead of forcing old code to use dynamic warpSize +// and losing constexpr. However, compiling for --offload-arch=amdgcnspirv +// could expose where C10_WARP_SIZE was used incorrectly where the dynamic +// warpSize is not allowed. +#if defined(__SPIRV__) +static __device__ inline int C10_WARP_SIZE_INTERNAL() { + return warpSize; +} +#else // __SPIRV__ +static __device__ inline constexpr int C10_WARP_SIZE_INTERNAL() { #if defined(__GFX9__) return 64; #else // __GFX9__ return 32; #endif // __GFX9__ } -#else // __HIPCC__ +#endif // __SPIRV__ +#if defined(__SPIRV__) +#define C10_WARP_SIZE_LOWER_BOUND 32 +#define C10_WARP_SIZE_UPPER_BOUND 64 +#elif defined(__GFX9__) +#define C10_WARP_SIZE_LOWER_BOUND 64 +#define C10_WARP_SIZE_UPPER_BOUND 64 +#else +#define C10_WARP_SIZE_LOWER_BOUND 32 +#define C10_WARP_SIZE_UPPER_BOUND 32 +#endif +#else // !__HIPCC__ static inline int C10_WARP_SIZE_INTERNAL() { return at::cuda::warp_size(); } +#define C10_WARP_SIZE_LOWER_BOUND 32 +#define C10_WARP_SIZE_UPPER_BOUND 64 #endif // __HIPCC__ - #define C10_WARP_SIZE (C10_WARP_SIZE_INTERNAL()) -#define C10_WARP_SIZE_STATIC 64 - -#else // defined(USE_ROCM) +#else // !USE_ROCM #define C10_WARP_SIZE 32 -#endif +#define C10_WARP_SIZE_LOWER_BOUND 32 +#define C10_WARP_SIZE_UPPER_BOUND 32 +#endif // USE_ROCM #if defined(_MSC_VER) && _MSC_VER <= 1900 #define __func__ __FUNCTION__ diff --git a/torch/headeronly/util/BFloat16.h b/torch/headeronly/util/BFloat16.h index 64479ba36f125..9aa08c265bd2c 100644 --- a/torch/headeronly/util/BFloat16.h +++ b/torch/headeronly/util/BFloat16.h @@ -12,7 +12,7 @@ #include #include -#if defined(__CUDACC__) && !defined(USE_ROCM) +#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702)) #include #endif @@ -46,7 +46,7 @@ struct alignas(2) BFloat16 { /* implicit */ inline C10_HOST_DEVICE BFloat16(float value); inline C10_HOST_DEVICE operator float() const; -#if defined(__CUDACC__) && !defined(USE_ROCM) +#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702)) inline C10_HOST_DEVICE BFloat16(const __nv_bfloat16& value); explicit inline C10_HOST_DEVICE operator __nv_bfloat16() const; #endif @@ -124,8 +124,9 @@ C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion") /// Constructors inline C10_HOST_DEVICE BFloat16::BFloat16(float value) : -#if defined(__CUDACC__) && !defined(USE_ROCM) && defined(__CUDA_ARCH__) && \ - __CUDA_ARCH__ >= 800 +#if defined(__CUDACC__) && \ + (!defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 || \ + defined(USE_ROCM) && (TORCH_HIP_VERSION >= 702)) x(__bfloat16_as_ushort(__float2bfloat16(value))) #elif defined(__SYCL_DEVICE_ONLY__) && \ defined(SYCL_EXT_ONEAPI_BFLOAT16_MATH_FUNCTIONS) @@ -139,7 +140,7 @@ inline C10_HOST_DEVICE BFloat16::BFloat16(float value) /// Implicit conversions inline C10_HOST_DEVICE BFloat16::operator float() const { -#if defined(__CUDACC__) && !defined(USE_ROCM) +#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702)) return __bfloat162float(*reinterpret_cast(&x)); #elif defined(__SYCL_DEVICE_ONLY__) && \ defined(SYCL_EXT_ONEAPI_BFLOAT16_MATH_FUNCTIONS) @@ -149,7 +150,7 @@ inline C10_HOST_DEVICE BFloat16::operator float() const { #endif } -#if defined(__CUDACC__) && !defined(USE_ROCM) +#if defined(__CUDACC__) && (!defined(USE_ROCM) || (TORCH_HIP_VERSION >= 702)) inline C10_HOST_DEVICE BFloat16::BFloat16(const __nv_bfloat16& value) { x = *reinterpret_cast(&value); } diff --git a/torch/headeronly/util/Float8_e4m3fn.h b/torch/headeronly/util/Float8_e4m3fn.h index b35fb1a7aa85d..06d70e6b4982c 100644 --- a/torch/headeronly/util/Float8_e4m3fn.h +++ b/torch/headeronly/util/Float8_e4m3fn.h @@ -49,6 +49,7 @@ struct alignas(1) Float8_e4m3fn { inline C10_HOST_DEVICE Float8_e4m3fn(float value); inline C10_HOST_DEVICE operator float() const; inline C10_HOST_DEVICE bool isnan() const; + inline C10_HOST_DEVICE bool isinf() const; }; inline std::ostream& operator<<(std::ostream& out, const Float8_e4m3fn& value) { @@ -256,6 +257,11 @@ inline C10_HOST_DEVICE bool Float8_e4m3fn::isnan() const { return (x & 0b01111111) == 0b01111111; } +inline C10_HOST_DEVICE bool Float8_e4m3fn::isinf() const { + // Note: fp8e4m3fn does not have infinity, so this always returns false. + return false; +} + /// Arithmetic inline C10_HOST_DEVICE Float8_e4m3fn diff --git a/torch/headeronly/util/Float8_e4m3fnuz.h b/torch/headeronly/util/Float8_e4m3fnuz.h index e361a2f92a2a5..4c11b3be05593 100644 --- a/torch/headeronly/util/Float8_e4m3fnuz.h +++ b/torch/headeronly/util/Float8_e4m3fnuz.h @@ -52,6 +52,7 @@ struct alignas(1) Float8_e4m3fnuz { inline C10_HOST_DEVICE Float8_e4m3fnuz(float value); inline C10_HOST_DEVICE operator float() const; inline C10_HOST_DEVICE bool isnan() const; + inline C10_HOST_DEVICE bool isinf() const; }; inline std::ostream& operator<<( @@ -160,6 +161,11 @@ inline C10_HOST_DEVICE bool Float8_e4m3fnuz::isnan() const { return x == 0b10000000; } +inline C10_HOST_DEVICE bool Float8_e4m3fnuz::isinf() const { + // Note: fp8e4m3fnuz does not have infinity, so this always returns false. + return false; +} + /// Arithmetic inline C10_HOST_DEVICE Float8_e4m3fnuz diff --git a/torch/lib/libshm/CMakeLists.txt b/torch/lib/libshm/CMakeLists.txt index c3cd26fea7bf3..688e7e7929147 100644 --- a/torch/lib/libshm/CMakeLists.txt +++ b/torch/lib/libshm/CMakeLists.txt @@ -21,7 +21,7 @@ target_include_directories(shm PUBLIC set_target_properties(shm PROPERTIES PREFIX "lib" IMPORT_PREFIX "lib" - CXX_STANDARD 17) + CXX_STANDARD 20) target_link_libraries(shm PRIVATE ${TORCH_CPU_LIB}) if(UNIX AND NOT APPLE) diff --git a/torch/lib/libshm/manager.cpp b/torch/lib/libshm/manager.cpp index 5d2c318c25142..3084f794154a9 100644 --- a/torch/lib/libshm/manager.cpp +++ b/torch/lib/libshm/manager.cpp @@ -46,12 +46,7 @@ static void register_fd(int fd) { } static void unregister_fd(int fd) { - pollfds.erase( - std::remove_if( - pollfds.begin(), - pollfds.end(), - [fd](const struct pollfd& pfd) { return pfd.fd == fd; }), - pollfds.end()); + std::erase_if(pollfds, [fd](const struct pollfd& pfd) { return pfd.fd == fd; }); client_sessions.erase(fd); } diff --git a/torch/nativert/executor/ConstantFolder.cpp b/torch/nativert/executor/ConstantFolder.cpp index bbae8d1d426ac..0cad2841e2391 100644 --- a/torch/nativert/executor/ConstantFolder.cpp +++ b/torch/nativert/executor/ConstantFolder.cpp @@ -102,7 +102,7 @@ void ConstantFolder::unlinkConstants( if (user == output) { continue; } - if (foldable.find(user) == foldable.end()) { + if (!foldable.contains(user)) { constFoldableCandidates.push(user); } } @@ -115,7 +115,7 @@ void ConstantFolder::unlinkConstants( // we only store folded values if there is a non-foldable user if (const auto& users = value->users(); std::any_of(users.begin(), users.end(), [&](const auto* u) { - return foldable.find(u) == foldable.end(); + return !foldable.contains(u); })) { foldedOutputValueIds_.insert(value->id()); } @@ -132,12 +132,7 @@ void ConstantFolder::unlinkConstants( // remove moved (i.e., associated w/ const-folded nodes) kernels // from the input kernel vector - kernels.erase( - std::remove_if( - kernels.begin(), - kernels.end(), - [](const auto& k) { return k == nullptr; }), - kernels.end()); + std::erase(kernels, nullptr); graph_.renumberValues(); graph_.finalize(); diff --git a/torch/nativert/executor/GraphExecutorBase.cpp b/torch/nativert/executor/GraphExecutorBase.cpp index a623d5873ea56..5cec883376e1f 100644 --- a/torch/nativert/executor/GraphExecutorBase.cpp +++ b/torch/nativert/executor/GraphExecutorBase.cpp @@ -79,6 +79,21 @@ ProfileMetrics GraphExecutorBase::benchmarkIndividualNodes( } } + // Capture input element counts per node (frame is populated after warmup). + results.inputElementsPerNode.resize(numNodes, 0); + for (const auto i : c10::irange(numNodes)) { + int64_t total_elements = 0; + for (const auto& input : nodeKernels_[i]->node()->inputs()) { + if (input.value && input.value->type().kind() == Type::Kind::Tensor) { + const auto& iv = executionFrame.getIValue(input.value->id()); + if (iv.isTensor()) { + total_elements += iv.toTensor().numel(); + } + } + } + results.inputElementsPerNode[i] = total_elements; + } + // Execute kernels caffe2::Timer timer; executionFrame.withManagedMemory([&](auto) { diff --git a/torch/nativert/executor/GraphExecutorBase.h b/torch/nativert/executor/GraphExecutorBase.h index dfe020ebae29e..d78d28253e201 100644 --- a/torch/nativert/executor/GraphExecutorBase.h +++ b/torch/nativert/executor/GraphExecutorBase.h @@ -21,6 +21,7 @@ struct ProfileMetrics { std::unordered_map instancesPerNodeType; std::unordered_set staticDispatchNodes; std::unordered_set primNodes; + std::vector inputElementsPerNode; float totalTime{0}; std::string name; }; diff --git a/torch/nativert/graph/Graph.cpp b/torch/nativert/graph/Graph.cpp index bb524d7d39bcf..3531f12c6a685 100644 --- a/torch/nativert/graph/Graph.cpp +++ b/torch/nativert/graph/Graph.cpp @@ -964,10 +964,7 @@ void Value::addUser(Node* node) { } void Value::eraseUser(Node* node) { - users_.erase( - std::remove_if( - users_.begin(), users_.end(), [&](Node* el) { return el == node; }), - users_.end()); + std::erase(users_, node); } std::vector Value::getListElements() const { @@ -1083,6 +1080,15 @@ std::ostream& operator<<(std::ostream& out, const Constant& constant) { out << fmt::format("{}", fmt::streamed(inner_list)); } out << ']'; + } else if constexpr (is_same_v>>) { + out << '['; + for (const auto& [idx, inner_list] : c10::enumerate(arg)) { + if (idx > 0) { + out << ", "; + } + out << fmt::format("{}", fmt::streamed(inner_list)); + } + out << ']'; } else if constexpr (is_same_v>) { out << fmt::format(""); VLOG(0) << "Subgraph pretty print is not implemented"; diff --git a/torch/nativert/graph/Graph.h b/torch/nativert/graph/Graph.h index ce5dd8e412e3d..6d8c449d15c14 100644 --- a/torch/nativert/graph/Graph.h +++ b/torch/nativert/graph/Graph.h @@ -100,6 +100,7 @@ using Constant = std::variant< std::vector, std::vector, std::vector>, + std::vector>, std::unique_ptr>; c10::IValue constantToIValue(const Constant& constant); diff --git a/torch/nativert/graph/GraphSignature.cpp b/torch/nativert/graph/GraphSignature.cpp index 7adff3b6a895a..84b4c10192564 100644 --- a/torch/nativert/graph/GraphSignature.cpp +++ b/torch/nativert/graph/GraphSignature.cpp @@ -258,6 +258,8 @@ GraphSignature::GraphSignature(const torch::_export::GraphSignature& storage) { userInputArg.tag() == torch::_export::Argument::Tag::AS_SYM_FLOATS || userInputArg.tag() == torch::_export::Argument::Tag::AS_INT_LISTS || + userInputArg.tag() == + torch::_export::Argument::Tag::AS_FLOAT_LISTS || userInputArg.tag() == torch::_export::Argument::Tag::AS_SCALAR_TYPE || userInputArg.tag() == diff --git a/torch/nativert/graph/Serialization.cpp b/torch/nativert/graph/Serialization.cpp index 98bcc34a0cbfd..ada34f15026dd 100644 --- a/torch/nativert/graph/Serialization.cpp +++ b/torch/nativert/graph/Serialization.cpp @@ -731,6 +731,17 @@ Constant constantToValue( } return ret; } + case torch::_export::Argument::Tag::AS_FLOAT_LISTS: { + std::vector> ret; + for (const auto& inner_list : jsonArg.get_as_float_lists()) { + std::vector inner_ret; + for (const auto& val : inner_list) { + inner_ret.push_back(val.get()); + } + ret.push_back(inner_ret); + } + return ret; + } case torch::_export::Argument::Tag::AS_STRING_TO_ARGUMENT: return None(); default: diff --git a/torch/nativert/python/Bindings.cpp b/torch/nativert/python/Bindings.cpp index d83c0bbc6ed43..de86306adea0a 100644 --- a/torch/nativert/python/Bindings.cpp +++ b/torch/nativert/python/Bindings.cpp @@ -3,6 +3,7 @@ #include #include +#ifdef FBCODE_CAFFE2 #include namespace py = pybind11; @@ -81,3 +82,22 @@ void initModelRunnerPybind(py::module& m) { } } // namespace torch::nativert + +#else // !FBCODE_CAFFE2 + +namespace py = pybind11; + +namespace torch::nativert { + +class StubModelRunner {}; + +// PyModelRunner is referenced from +// https://github.com/pytorch/benchmark/blob/b8d35ba51a3149b7212888b4010ddee97f19947f/userbenchmark/dynamo/dynamobench/common.py#L45 +void initModelRunnerPybind(py::module& m) { + py::class_>( + m, "PyModelRunner"); +} + +} // namespace torch::nativert + +#endif // FBCODE_CAFFE2 diff --git a/torch/nn/attention/flex_attention.py b/torch/nn/attention/flex_attention.py index 7365e726b055c..5cf2c42e2b93c 100644 --- a/torch/nn/attention/flex_attention.py +++ b/torch/nn/attention/flex_attention.py @@ -9,6 +9,7 @@ import itertools import math import operator +import types import typing import warnings from collections.abc import Callable @@ -21,11 +22,18 @@ from torch._higher_order_ops.flex_attention import flex_attention as flex_attention_hop from torch._higher_order_ops.utils import setup_compilation_env from torch.nn.attention._utils import _validate_sdpa_input -from torch.utils._pytree import GetAttrKey, tree_map_only +from torch.utils._pytree import ( + GetAttrKey, + tree_flatten, + tree_map_only, + tree_unflatten, + TreeSpec, +) if typing.TYPE_CHECKING: from torch._prims_common import DeviceLikeType + from torch.fx.node import BaseArgumentTypes # Private debug flag to disable internal compilation wrapping for debugging purposes. @@ -442,49 +450,228 @@ def _adjust_num_blocks_and_indices( return num_blocks, indices -def _closure_contents(fn: object) -> tuple[object, ...]: - """Extract closure cell contents for comparison.""" - closure = getattr(fn, "__closure__", None) - if closure is None: - return () - return tuple(cell.cell_contents for cell in closure) +# TreeSpec for an empty tuple — used as the sentinel when there are no closure leaves. +_EMPTY_CLOSURE_SPEC = tree_flatten(())[1] + + +class _ExtractedLeaf: + """Sentinel in _StrippedClosure.leaf_entries marking a position that is + filled from the extracted pytree leaves list during reconstruction.""" + + __slots__ = () + + def __repr__(self) -> str: + return "_EXTRACTED_LEAF" + + +_EXTRACTED_LEAF = _ExtractedLeaf() + + +class _FunctionLeaf(typing.NamedTuple): + """Entry in _StrippedClosure.leaf_entries for a recursively processed + function. Stores enough information to reconstruct the function from + the extracted leaves during unflattening.""" + + stripped: _StrippedClosure | Callable[..., Any] + closure_spec: TreeSpec + n_extracted: int # number of extracted leaves this function contributes + + +class _StrippedClosure(typing.NamedTuple): + """Data container holding the parts of a function needed for reconstruction. + + Created by _extract_closure_pytree when closure tensors are lifted into + pytree leaves. Unlike a FunctionType with None-filled cells, this is not + callable — it is pure data stored in the pytree context. + """ + + code: types.CodeType + globals_dict: dict[str, Any] + name: str + qualname: str + defaults: tuple[Any, ...] | None + kwdefaults: dict[str, Any] | None + extra_dict: dict[str, Any] + # Per-position info for the closure's flattened leaves. + # _EXTRACTED_LEAF → position filled from the extracted leaves list. + # _FunctionLeaf → recursively processed function (not a valid pytree leaf). + leaf_entries: tuple[_ExtractedLeaf | _FunctionLeaf, ...] + + +def _extract_closure_pytree( + fn, _seen: set[int] | None = None +) -> tuple[ + tuple[BaseArgumentTypes, ...], TreeSpec, _StrippedClosure | Callable[..., Any] +]: + """Extract closure contents as a flattened sub-pytree. + + Returns (extracted_leaves, closure_spec, fn_or_stripped) where: + - extracted_leaves: flattened non-function contents from the closure, + plus any tensors/scalars recursively extracted from nested function + closures + - closure_spec: TreeSpec describing how to reconstruct the closure contents + - fn_or_stripped: either the original fn (no extraction) or a + _StrippedClosure carrying the function parts needed for reconstruction + + Functions found among the closure leaves are recursively processed: their + own closure tensors are extracted into the leaves list, and their skeleton + is stored in _StrippedClosure.leaf_entries as a _FunctionLeaf. All other + values (tensors, scalars, None, etc.) remain as extracted leaves. + + If fn is not a plain function, has no closure, or has empty cells, returns + the original function unchanged with no closure leaves. + + Skipped under Dynamo tracing (torch.compiler.is_compiling) because Dynamo + can't trace through closure cell introspection and handles freevars via its + own lifting mechanism. + """ + if not inspect.isfunction(fn) or torch.compiler.is_compiling(): + return (), _EMPTY_CLOSURE_SPEC, fn + + # Cycle detection for self-referencing closures. + if _seen is None: + _seen = set() + if id(fn) in _seen: + return (), _EMPTY_CLOSURE_SPEC, fn + _seen.add(id(fn)) + + closure = fn.__closure__ + if not closure: + return (), _EMPTY_CLOSURE_SPEC, fn + + try: + contents = tuple(cell.cell_contents for cell in closure) + except ValueError: + # Empty cell (created but not yet assigned) — can't extract + return (), _EMPTY_CLOSURE_SPEC, fn + + closure_leaves, closure_spec = tree_flatten(contents) + + extracted: list[BaseArgumentTypes] = [] + leaf_entries: list[_ExtractedLeaf | _FunctionLeaf] = [] + for leaf in closure_leaves: + if inspect.isfunction(leaf): + child_extracted, child_spec, child_stripped = _extract_closure_pytree( + leaf, _seen + ) + extracted.extend(child_extracted) + leaf_entries.append( + _FunctionLeaf(child_stripped, child_spec, len(child_extracted)) + ) + else: + extracted.append(leaf) + leaf_entries.append(_EXTRACTED_LEAF) + + stripped = _StrippedClosure( + code=fn.__code__, + globals_dict=fn.__globals__, + name=fn.__name__, + qualname=fn.__qualname__, + defaults=fn.__defaults__, + kwdefaults=fn.__kwdefaults__, + extra_dict=dict(fn.__dict__) if fn.__dict__ else {}, + leaf_entries=tuple(leaf_entries), + ) + + return tuple(extracted), closure_spec, stripped + + +def _reconstruct_closure_fn(stripped, extracted_leaves, closure_spec): + """Rebuild a function from a _StrippedClosure and flattened extracted leaves.""" + if not isinstance(stripped, _StrippedClosure): + return stripped + + all_leaves: list[BaseArgumentTypes | Callable[..., Any]] = [] + idx = 0 + for entry in stripped.leaf_entries: + if isinstance(entry, _FunctionLeaf): + child_fn = _reconstruct_closure_fn( + entry.stripped, + extracted_leaves[idx : idx + entry.n_extracted], + entry.closure_spec, + ) + all_leaves.append(child_fn) + idx += entry.n_extracted + else: + # _EXTRACTED_LEAF — take from extracted leaves + all_leaves.append(extracted_leaves[idx]) + idx += 1 + + contents = tree_unflatten(all_leaves, closure_spec) + new_cells = tuple(types.CellType(v) for v in contents) + + restored = types.FunctionType( + stripped.code, + stripped.globals_dict, + stripped.name, + stripped.defaults, + new_cells, + ) + restored.__qualname__ = stripped.qualname + if stripped.kwdefaults: + restored.__kwdefaults__ = stripped.kwdefaults + if stripped.extra_dict: + restored.__dict__.update(stripped.extra_dict) + + return restored class _MaskModWrapper: - """Wraps a mask_mod function with value-based equality. + """Wraps a mask_mod or _StrippedClosure with value-based equality. BlockMask stores an arbitrary callable (mask_mod) in its pytree context. The default __eq__ for functions uses identity comparison, which is too strict when the same closure is recreated (e.g., defined inside forward()). - This wrapper compares functions by their code object and closure contents. + + When closure tensors have been extracted (by _extract_closure_pytree), fn + is a _StrippedClosure (pure data, not callable). Equality compares the + code objects + closure_spec — no tensor dispatch is triggered. + + When extraction is skipped (e.g., under Dynamo), fn is the original + callable and equality compares code objects + closure contents (for plain + functions) or delegates to __eq__ (for callable objects). """ - __slots__ = ("fn",) + __slots__ = ("fn", "closure_spec") - def __init__(self, fn: _mask_mod_signature) -> None: + def __init__(self, fn, closure_spec=None) -> None: self.fn = fn + self.closure_spec = closure_spec def __call__(self, b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor: + if isinstance(self.fn, _StrippedClosure): + raise RuntimeError( + "_MaskModWrapper with _StrippedClosure is not callable — " + "use _reconstruct_closure_fn to rebuild the function first" + ) return self.fn(b, h, q_idx, kv_idx) def __eq__(self, other: object) -> bool: if not isinstance(other, _MaskModWrapper): return False - if self.fn is other.fn: + if self.fn is other.fn and self.closure_spec is other.closure_spec: return True - if ( - inspect.isfunction(self.fn) - and inspect.isfunction(other.fn) - and self.fn.__code__ == other.fn.__code__ - and _closure_contents(self.fn) == _closure_contents(other.fn) + # Extracted case: _StrippedClosure — compare code + closure_spec + if isinstance(self.fn, _StrippedClosure) and isinstance( + other.fn, _StrippedClosure + ): + return ( + self.fn.code == other.fn.code + and self.closure_spec == other.closure_spec + ) + # Non-extracted plain functions: compare code + closure contents + if inspect.isfunction(self.fn) and inspect.isfunction(other.fn): + return self.fn.__code__ == other.fn.__code__ + # Callable objects: delegate to their __eq__ + if not isinstance(self.fn, _StrippedClosure) and not isinstance( + other.fn, _StrippedClosure ): - return True - # For callable objects (not plain functions), delegate to their __eq__ - if not inspect.isfunction(self.fn) and not inspect.isfunction(other.fn): return self.fn == other.fn return False def __hash__(self) -> int: + if isinstance(self.fn, _StrippedClosure): + return hash(self.fn.code) if inspect.isfunction(self.fn): return hash(self.fn.__code__) return hash(self.fn) @@ -1043,30 +1230,43 @@ def _unwrap_context_value(attr: str, value: Any) -> Any: def _flatten( self, - ) -> tuple[tuple[Tensor | None, ...], tuple[Any, ...]]: + ) -> tuple[tuple[BaseArgumentTypes | None, ...], tuple[Any, ...]]: """Flatten BlockMask into a list of tensors and context. - Wraps mask_mod in _MaskModWrapper for value-based comparison in TreeSpec. + Closure tensors from mask_mod are extracted into the leaves via + _extract_closure_pytree so they are visible to the tracing + infrastructure (instead of being hidden in the pytree context). """ tensors = tuple(getattr(self, attr) for attr in self._TENSOR_ATTRS) + closure_leaves, closure_spec, stripped = _extract_closure_pytree(self.mask_mod) + all_leaves = tensors + closure_leaves context = tuple( self._wrap_context_value(attr, getattr(self, attr)) + if attr != "mask_mod" + else _MaskModWrapper(stripped, closure_spec) for attr in self._CONTEXT_ATTRS ) - return tensors, context + return all_leaves, context @classmethod def _unflatten( cls, - tensors: tuple[Tensor | None, ...], + leaves: tuple[Any, ...], context: tuple[Any, ...], ) -> Self: - """Unflatten tensors and context back into a BlockMask.""" - kwargs = { - attr: cls._unwrap_context_value(attr, val) - for attr, val in zip(cls._CONTEXT_ATTRS, context) - } - kwargs.update(zip(cls._TENSOR_ATTRS, tensors)) + """Unflatten leaves and context back into a BlockMask.""" + n_regular = len(cls._TENSOR_ATTRS) + regular_leaves = leaves[:n_regular] + closure_leaves = leaves[n_regular:] + kwargs = {} + for attr, val in zip(cls._CONTEXT_ATTRS, context): + if attr == "mask_mod" and isinstance(val, _MaskModWrapper): + kwargs[attr] = _reconstruct_closure_fn( + val.fn, closure_leaves, val.closure_spec + ) + else: + kwargs[attr] = cls._unwrap_context_value(attr, val) + kwargs.update(zip(cls._TENSOR_ATTRS, regular_leaves)) return cls(**kwargs) def _flatten_with_keys( @@ -1074,16 +1274,25 @@ def _flatten_with_keys( ) -> tuple[tuple[tuple[GetAttrKey, Any], ...], tuple[tuple[GetAttrKey, Any], ...]]: """Flatten BlockMask with keys for better tracing. - Wraps mask_mod in _MaskModWrapper for value-based comparison in TreeSpec. + Closure tensors from mask_mod are extracted into the leaves via + _extract_closure_pytree so they are visible to the tracing + infrastructure (instead of being hidden in the pytree context). """ tensors = tuple( (GetAttrKey(attr), getattr(self, attr)) for attr in self._TENSOR_ATTRS ) + closure_leaves, closure_spec, stripped = _extract_closure_pytree(self.mask_mod) + closure_with_keys = tuple( + (GetAttrKey(f"_closure_{i}"), leaf) for i, leaf in enumerate(closure_leaves) + ) + all_leaves = tensors + closure_with_keys context = tuple( (GetAttrKey(attr), self._wrap_context_value(attr, getattr(self, attr))) + if attr != "mask_mod" + else (GetAttrKey(attr), _MaskModWrapper(stripped, closure_spec)) for attr in self._CONTEXT_ATTRS ) - return tensors, context + return all_leaves, context def _broadcast_to_dim(x: Tensor, dim: int) -> Tensor: diff --git a/torch/nn/attention/varlen.py b/torch/nn/attention/varlen.py index 4401754e9395b..96cfc332b8699 100644 --- a/torch/nn/attention/varlen.py +++ b/torch/nn/attention/varlen.py @@ -54,6 +54,7 @@ def _varlen_attn( is_causal: bool = False, scale: float | None = None, window_size: list[int] | None = None, + enable_gqa: bool = False, seqused_k: torch.Tensor | None = None, block_table: torch.Tensor | None = None, num_splits: int | None = None, @@ -70,6 +71,9 @@ def _varlen_attn( if use_cudnn: log.info("Using cuDNN backend for varlen_attn") + if enable_gqa: + # TODO: check this + raise RuntimeError("GQA is not supported with the cuDNN backend.") if num_splits is not None: # TODO: check this raise RuntimeError("num_splits is not supported with the cuDNN backend.") @@ -140,6 +144,7 @@ def _varlen_attn_fake( is_causal: bool = False, scale: float | None = None, window_size: list[int] | None = None, + enable_gqa: bool = False, seqused_k: torch.Tensor | None = None, block_table: torch.Tensor | None = None, num_splits: int | None = None, @@ -159,16 +164,18 @@ def _varlen_attn_fake( # For varlen path: logsumexp shape is (num_heads, total_q) total_q = query.size(0) num_heads = query.size(1) + logsumexp = torch.empty( + (num_heads, total_q), dtype=torch.float, device=query.device + ) + if torch.version.hip: - # ROCm uses batched format: [batch_size, num_heads, max_q] - batch_size = cu_seq_q.size(0) - 1 - logsumexp = torch.empty( - (batch_size, num_heads, max_q), dtype=torch.float, device=query.device - ) - else: - logsumexp = torch.empty( - (num_heads, total_q), dtype=torch.float, device=query.device - ) + preferred = torch._C._get_rocm_fa_preferred_backend() + if preferred == torch._C._ROCmFABackend.AOTriton: + # AOTriton ROCm path uses batched 3D + batch_size = cu_seq_q.size(0) - 1 + logsumexp = torch.empty( + (batch_size, num_heads, max_q), dtype=torch.float, device=query.device + ) rng_state = torch.empty((2,), dtype=torch.uint64, device=query.device) @@ -187,6 +194,7 @@ def varlen_attn( return_aux: AuxRequest | None = None, scale: float | None = None, window_size: tuple[int, int] = (-1, -1), + enable_gqa: bool = False, seqused_k: torch.Tensor | None = None, block_table: torch.Tensor | None = None, num_splits: int | None = None, @@ -197,11 +205,11 @@ def varlen_attn( variable-length sequences using cumulative sequence position tensors. Args: - query (Tensor): Query tensor; shape :math:`(T_q, H, D)` - key (Tensor): Key tensor; shape :math:`(T_k, H, D)`, or - :math:`(\text{total\_pages}, \text{page\_size}, H, D)` when ``block_table`` is provided. - value (Tensor): Value tensor; shape :math:`(T_k, H, D)`, or - :math:`(\text{total\_pages}, \text{page\_size}, H, D)` when ``block_table`` is provided. + query (Tensor): Query tensor; shape :math:`(T_q, H_q, D)` + key (Tensor): Key tensor; shape :math:`(T_k, H_{kv}, D)`, or + :math:`(\text{total\_pages}, \text{page\_size}, H_{kv}, D)` when ``block_table`` is provided. + value (Tensor): Value tensor; shape :math:`(T_k, H_{kv}, D)`, or + :math:`(\text{total\_pages}, \text{page\_size}, H_{kv}, D)` when ``block_table`` is provided. cu_seq_q (Tensor): Cumulative sequence positions for queries; shape :math:`(N+1,)` cu_seq_k (Tensor): Cumulative sequence positions for keys/values; shape :math:`(N+1,)` max_q (int): Maximum query sequence length in the batch. @@ -211,6 +219,11 @@ def varlen_attn( window_size (tuple[int, int], optional): Window size for sliding window attention as (left, right). Use (-1, -1) for full attention (default), (-1, 0) for causal attention, or (W, 0) for causal attention with sliding window of size W. + enable_gqa (bool): If set to True, enables Grouped Query Attention (GQA) + and allows key/value to have fewer heads than query. + Each KV head is shared by a group of :math:`H_q / H_{kv}` query heads, + so :math:`H_q` must be divisible by :math:`H_{kv}`. + Default is False. seqused_k (Tensor, optional): Number of valid KV tokens per batch element; shape :math:`(N,)`. When set, only the first ``seqused_k[i]`` tokens in the key/value sequence for batch element *i* participate in attention. Useful for KV-cache decoding where the cache slot @@ -239,16 +252,17 @@ def varlen_attn( ``None`` (default), the kernel chooses automatically. Returns: - output (Tensor): Output tensor from attention computation; shape :math:`(T_q, H, D)`. + output (Tensor): Output tensor from attention computation; shape :math:`(T_q, H_q, D)`. If ``return_aux`` is not None and ``return_aux.lse`` is True: - lse (Tensor): Log-sum-exp of attention scores; shape :math:`(T_q, H)`. + lse (Tensor): Log-sum-exp of attention scores; shape :math:`(T_q, H_q)`. Shape legend: - :math:`N`: Batch size - :math:`T_q`: Total number of query tokens in the batch (sum of all query sequence lengths) - :math:`T_k`: Total number of key/value tokens in the batch (sum of all key/value sequence lengths) - - :math:`H`: Number of attention heads + - :math:`H_q`: Number of query attention heads + - :math:`H_{kv}`: Number of key/value attention heads (equal to :math:`H_q` unless GQA is enabled) - :math:`D`: Head dimension Example:: @@ -285,6 +299,20 @@ def varlen_attn( ... ) """ + num_heads_q = query.size(1) + num_heads_k = key.size(2) if block_table is not None else key.size(1) + if not enable_gqa and num_heads_q != num_heads_k: + raise ValueError( + f"Expect query and key/value to have the same number of heads " + f"but got Hq={num_heads_q} and Hkv={num_heads_k}. " + f"Try setting enable_gqa=True for GQA." + ) + if enable_gqa and num_heads_q % num_heads_k != 0: + raise ValueError( + f"Expect number of query heads to be a multiple of kv heads for GQA " + f"but got Hq={num_heads_q} and Hkv={num_heads_k}." + ) + is_causal = window_size == (-1, 0) out, lse, _ = torch.ops.torch_attn._varlen_attn( query, @@ -297,6 +325,7 @@ def varlen_attn( is_causal, scale, list(window_size), + enable_gqa, seqused_k, block_table, num_splits, @@ -319,6 +348,7 @@ def _varlen_attn_out( is_causal: bool = False, scale: float | None = None, window_size: list[int] | None = None, + enable_gqa: bool = False, seqused_k: torch.Tensor | None = None, block_table: torch.Tensor | None = None, num_splits: int | None = None, @@ -372,6 +402,7 @@ def _varlen_attn_out_fake( is_causal: bool = False, scale: float | None = None, window_size: list[int] | None = None, + enable_gqa: bool = False, seqused_k: torch.Tensor | None = None, block_table: torch.Tensor | None = None, num_splits: int | None = None, @@ -409,6 +440,7 @@ def varlen_attn_out( return_aux: AuxRequest | None = None, scale: float | None = None, window_size: tuple[int, int] = (-1, -1), + enable_gqa: bool = False, seqused_k: torch.Tensor | None = None, block_table: torch.Tensor | None = None, num_splits: int | None = None, @@ -419,6 +451,20 @@ def varlen_attn_out( instead of allocating a new one. """ + num_heads_q = query.size(1) + num_heads_k = key.size(2) if block_table is not None else key.size(1) + if not enable_gqa and num_heads_q != num_heads_k: + raise ValueError( + f"Expect query and key/value to have the same number of heads " + f"but got Hq={num_heads_q} and Hkv={num_heads_k}. " + f"Try setting enable_gqa=True for GQA." + ) + if enable_gqa and num_heads_q % num_heads_k != 0: + raise ValueError( + f"Expect number of query heads to be a multiple of kv heads for GQA " + f"but got Hq={num_heads_q} and Hkv={num_heads_k}." + ) + is_causal = window_size == (-1, 0) lse = torch.ops.torch_attn._varlen_attn_out( out, @@ -432,6 +478,7 @@ def varlen_attn_out( is_causal, scale, list(window_size), + enable_gqa, seqused_k, block_table, num_splits, @@ -453,6 +500,7 @@ def _setup_context(ctx: Any, inputs: tuple[Any, ...], output: Any) -> None: is_causal, scale, window_size, + enable_gqa, seqused_k, block_table, num_splits, @@ -598,7 +646,9 @@ def _backward( scale, window_size, ) - num_params = 10 # cu_seq_q, cu_seq_k, max_q, max_k, is_causal, scale, window_size, seqused_k, block_table, num_splits + # cu_seq_q, cu_seq_k, max_q, max_k, is_causal, scale, window_size, \ + # enable_gqa, seqused_k, block_table, num_splits + num_params = 11 return (dq, dk, dv, *((None,) * num_params)) @@ -607,3 +657,15 @@ def _backward( torch._dynamo.disallow_in_graph( torch.ops.aten._flash_attention_forward_no_dropout_inplace ) + +from torch.utils.flop_counter import ( + _varlen_attn_backward_flop, + _varlen_attn_forward_flop, + _varlen_attn_out_flop, + flop_registry, +) + + +flop_registry[torch.ops.torch_attn._varlen_attn] = _varlen_attn_forward_flop +flop_registry[torch.ops.torch_attn._varlen_attn_out] = _varlen_attn_out_flop +flop_registry[torch.ops.torch_attn._varlen_attn_backward] = _varlen_attn_backward_flop diff --git a/torch/nn/functional.py b/torch/nn/functional.py index 88fa9149b21c4..0906e86ee623c 100644 --- a/torch/nn/functional.py +++ b/torch/nn/functional.py @@ -4497,7 +4497,7 @@ def upsample( # noqa: F811 `mini-batch x channels x [optional depth] x [optional height] x width`. The modes available for upsampling are: `nearest`, `linear` (3D-only), - `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only) + `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only), `lanczos` (4D-only) Args: input (Tensor): the input tensor @@ -4506,7 +4506,7 @@ def upsample( # noqa: F811 scale_factor (float or Tuple[float]): multiplier for spatial size. Has to match input size if it is a tuple. mode (str): algorithm used for upsampling: ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` | - ``'trilinear'``. Default: ``'nearest'`` + ``'trilinear'`` | ``'lanczos'``. Default: ``'nearest'`` align_corners (bool, optional): Geometrically, we consider the pixels of the input and output as squares rather than points. If set to ``True``, the input and output tensors are aligned by the @@ -4519,7 +4519,7 @@ def upsample( # noqa: F811 Default: ``False`` .. note:: - With ``mode='bicubic'``, it's possible to cause overshoot, in other words it can produce + With ``mode='bicubic'`` or ``mode='lanczos'``, it's possible to cause overshoot, in other words it can produce negative values or values greater than 255 for images. Explicitly call ``result.clamp(min=0, max=255)`` if you want to reduce the overshoot when displaying the image. @@ -4637,7 +4637,7 @@ def interpolate( # noqa: F811 `mini-batch x channels x [optional depth] x [optional height] x width`. The modes available for resizing are: `nearest`, `linear` (3D-only), - `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only), `area`, `nearest-exact` + `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only), `lanczos` (4D-only, CPU only), `area`, `nearest-exact` Args: input (Tensor): the input tensor @@ -4647,7 +4647,7 @@ def interpolate( # noqa: F811 its length has to match the number of spatial dimensions; `input.dim() - 2`. mode (str): algorithm used for upsampling: ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` | - ``'trilinear'`` | ``'area'`` | ``'nearest-exact'``. Default: ``'nearest'`` + ``'trilinear'`` | ``'lanczos'`` | ``'area'`` | ``'nearest-exact'``. Default: ``'nearest'`` align_corners (bool, optional): Geometrically, we consider the pixels of the input and output as squares rather than points. If set to ``True``, the input and output tensors are aligned by the @@ -4669,14 +4669,19 @@ def interpolate( # noqa: F811 be used directly for interpolation. Default: ``None``. antialias (bool, optional): flag to apply anti-aliasing. Default: ``False``. Using anti-alias option together with ``align_corners=False``, interpolation result would match Pillow - result for downsampling operation. Supported modes: ``'bilinear'``, ``'bicubic'``. + result for downsampling operation. Supported modes: ``'bilinear'``, ``'bicubic'``, ``'lanczos'``. .. note:: - With ``mode='bicubic'``, it's possible to cause overshoot. For some dtypes, it can produce + With ``mode='bicubic'`` or ``mode='lanczos'``, it's possible to cause overshoot. For some dtypes, it can produce negative values or values greater than 255 for images. Explicitly call ``result.clamp(min=0,max=255)`` if you want to reduce the overshoot when displaying the image. For ``uint8`` inputs, it already performs saturating cast operation. So, no manual `clamp` operation is needed. + .. note:: + Mode ``mode='lanczos'`` uses a Lanczos-3 windowed sinc filter (6 taps) and requires + ``antialias=True``. It only supports 4-D input (i.e. 2D spatial) and CPU. With ``antialias=True`` + and ``align_corners=False``, the result matches PIL's ``Image.LANCZOS`` resampling filter. + .. note:: Mode ``mode='nearest-exact'`` matches Scikit-Image and PIL nearest neighbours interpolation algorithms and fixes known issues with ``mode='nearest'``. This mode is introduced to keep @@ -4807,9 +4812,11 @@ def interpolate( # noqa: F811 ] scale_factors = None - if antialias and not (mode in ("bilinear", "bicubic") and input.ndim == 4): + if antialias and not ( + mode in ("bilinear", "bicubic", "lanczos") and input.ndim == 4 + ): raise ValueError( - "Anti-alias option is restricted to bilinear and bicubic modes and requires a 4-D tensor as input" + "Anti-alias option is restricted to bilinear, bicubic, and lanczos modes and requires a 4-D tensor as input" ) if input.dim() == 3 and mode == "nearest": @@ -4923,6 +4930,21 @@ def interpolate( # noqa: F811 scale_factors, ) + if input.dim() == 4 and mode == "lanczos": + if align_corners is None: + raise AssertionError("align_corners is unexpectedly None") + if align_corners: + raise ValueError("Lanczos mode does not support align_corners=True") + if not antialias: + raise ValueError("Lanczos mode requires antialias=True") + return torch._C._nn._upsample_lanczos2d_aa( + input, + # pyrefly: ignore [bad-argument-type] + output_size, + align_corners, + scale_factors, + ) + if input.dim() == 3 and mode == "bilinear": raise NotImplementedError("Got 3D input, but bilinear mode needs 4D input") if input.dim() == 3 and mode == "trilinear": @@ -4938,7 +4960,7 @@ def interpolate( # noqa: F811 raise NotImplementedError( "Input Error: Only 3D, 4D and 5D input Tensors supported" - f" (got {input.dim()}D) for the modes: nearest | linear | bilinear | bicubic | trilinear | area | nearest-exact" + f" (got {input.dim()}D) for the modes: nearest | linear | bilinear | bicubic | trilinear | lanczos | area | nearest-exact" f" (got {mode})" ) @@ -5974,7 +5996,7 @@ def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0. attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device) if is_causal: assert attn_mask is None - temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) + temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(diagonal=0) attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) if attn_mask is not None: @@ -6078,7 +6100,7 @@ def forward(self, ...): key (Tensor): Key tensor; shape :math:`(N, ..., H, S, E)`. value (Tensor): Value tensor; shape :math:`(N, ..., H, S, Ev)`. attn_mask (optional Tensor): Attention mask; shape must be broadcastable to the shape of attention weights, - which is :math:`(N,..., L, S)`. Two types of masks are supported. + which is :math:`(N,..., Hq, L, S)`. Two types of masks are supported. A boolean mask where a value of True indicates that the element *should* take part in attention. A float mask of the same type as query, key, value that is added to the attention score. dropout_p (float): Dropout probability; if greater than 0.0, dropout is applied @@ -6560,10 +6582,10 @@ def multi_head_attention_forward( # # reshape q, k, v for multihead attention and make them batch first # - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1) if static_k is None: - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1) else: # TODO finish disentangling control flow so we don't do in-projections when statics are passed @@ -6577,7 +6599,7 @@ def multi_head_attention_forward( ) k = static_k if static_v is None: - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1) else: # TODO finish disentangling control flow so we don't do in-projections when statics are passed @@ -6650,11 +6672,15 @@ def multi_head_attention_forward( ) else: attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1)) + if not torch.jit.is_scripting(): + del q_scaled, k attn_output_weights = softmax(attn_output_weights, dim=-1) if dropout_p > 0.0: attn_output_weights = dropout(attn_output_weights, p=dropout_p) attn_output = torch.bmm(attn_output_weights, v) + if not torch.jit.is_scripting(): + del v attn_output = ( attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim) @@ -6682,15 +6708,22 @@ def multi_head_attention_forward( else: attn_mask = attn_mask.view(bsz, num_heads, -1, src_len) + # pyrefly: ignore [bad-argument-type] q = q.view(bsz, num_heads, tgt_len, head_dim) - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type] k = k.view(bsz, num_heads, src_len, head_dim) - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type] v = v.view(bsz, num_heads, src_len, head_dim) attn_output = scaled_dot_product_attention( q, k, v, attn_mask, dropout_p, is_causal ) + # Free q, k, v and their backing projection storage before the + # .contiguous() call below allocates. In self-attention the three + # tensors are views of a single packed projection, so releasing all + # references here lets the allocator reclaim that memory immediately. + if not torch.jit.is_scripting(): + del q, k, v attn_output = ( attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim) ) diff --git a/torch/nn/modules/_functions.py b/torch/nn/modules/_functions.py index 408e6ef42f128..59352264b3216 100644 --- a/torch/nn/modules/_functions.py +++ b/torch/nn/modules/_functions.py @@ -124,6 +124,7 @@ def forward( return torch.empty_like(input) @staticmethod + # pyrefly: ignore [bad-override] def backward(self, grad_output): if not ( grad_output.is_contiguous(memory_format=torch.channels_last) diff --git a/torch/nn/modules/activation.py b/torch/nn/modules/activation.py index 658dd0a4d4fa3..739b2a95fc3c1 100644 --- a/torch/nn/modules/activation.py +++ b/torch/nn/modules/activation.py @@ -1394,6 +1394,7 @@ def forward( elif torch.is_autocast_enabled(): why_not_fast_path = "autocast is enabled" + fast_path_blocked_by_tracing = False if not why_not_fast_path: tensor_args = ( query, @@ -1408,8 +1409,6 @@ def forward( # generator expressions. if torch.overrides.has_torch_function(tensor_args): why_not_fast_path = "some Tensor argument has_torch_function" - elif _is_make_fx_tracing(): - why_not_fast_path = "we are running make_fx tracing" elif not all(_check_arg_device(x) for x in tensor_args): why_not_fast_path = ( "some Tensor argument's device is neither one of " @@ -1422,6 +1421,9 @@ def forward( "grad is enabled and at least one of query or the " "input/output projection weights or biases requires_grad" ) + elif _is_make_fx_tracing(): + why_not_fast_path = "we are running make_fx tracing" + fast_path_blocked_by_tracing = True if not why_not_fast_path: merged_mask, mask_type = self.merge_masks( attn_mask, key_padding_mask, query @@ -1511,7 +1513,11 @@ def forward( is_causal=is_causal, ) if self.batch_first and is_batched: - return attn_output.transpose(1, 0), attn_output_weights + attn_output = attn_output.transpose(1, 0) + if fast_path_blocked_by_tracing: + # Keep the traced slowpath layout aligned with eager fastpath. + attn_output = attn_output.contiguous() + return attn_output, attn_output_weights else: return attn_output, attn_output_weights diff --git a/torch/nn/modules/conv.py b/torch/nn/modules/conv.py index 88b23084442a2..bcc39e3e24f05 100644 --- a/torch/nn/modules/conv.py +++ b/torch/nn/modules/conv.py @@ -400,9 +400,11 @@ class Conv2d(_ConvNd): where :math:`\star` is the valid 2D `cross-correlation`_ operator, - :math:`N` is a batch size, :math:`C` denotes a number of channels, - :math:`H` is a height of input planes in pixels, and :math:`W` is - width in pixels. + :math:`N` is a batch size, :math:`C_{\text{in}}` and :math:`C_{\text{out}}` correspond to + :attr:`in_channels` and :attr:`out_channels` respectively, + :math:`H` and :math:`W` are the input height and width in pixels. + See the Shape section below for how :math:`H_{\text{out}}` and :math:`W_{\text{out}}` + are derived from :attr:`kernel_size`, :attr:`stride`, :attr:`padding`, and :attr:`dilation`. """ + r""" diff --git a/torch/nn/modules/normalization.py b/torch/nn/modules/normalization.py index cab31782edbde..c4df21528cae2 100644 --- a/torch/nn/modules/normalization.py +++ b/torch/nn/modules/normalization.py @@ -365,10 +365,10 @@ class RMSNorm(Module): If a single integer is used, it is treated as a singleton list, and this module will normalize over the last dimension which is expected to be of that specific size. - eps: a value added to the denominator for numerical stability. If not specified, - uses the machine epsilon of the computation (opmath) type: fp16/bf16 and - fp32 inputs use ``torch.finfo(torch.float32).eps``, while fp64 inputs use - ``torch.finfo(torch.float64).eps``. + eps (float, optional): a value added to the denominator for numerical stability. + If not specified, uses the machine epsilon of the computation (opmath) type: + fp16/bf16 and fp32 inputs use ``torch.finfo(torch.float32).eps``, while fp64 + inputs use ``torch.finfo(torch.float64).eps``. Default: ``None`` elementwise_affine: a boolean value that when set to ``True``, this module has learnable per-element affine parameters initialized to ones (for weights). Default: ``True``. diff --git a/torch/nn/parallel/data_parallel.py b/torch/nn/parallel/data_parallel.py index 65e78dfe3180a..16537624663b1 100644 --- a/torch/nn/parallel/data_parallel.py +++ b/torch/nn/parallel/data_parallel.py @@ -259,7 +259,7 @@ def data_parallel( device_ids = [_get_device_index(x, True) for x in device_ids] output_device = _get_device_index(output_device, True) - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] src_device_obj = torch.device(device_type, device_ids[0]) # pyrefly: ignore [bad-argument-type] diff --git a/torch/nn/parallel/distributed.py b/torch/nn/parallel/distributed.py index b7fa3865fa4fe..50f230b995fba 100644 --- a/torch/nn/parallel/distributed.py +++ b/torch/nn/parallel/distributed.py @@ -654,7 +654,7 @@ class DistributedDataParallel(Module, Joinable): If you plan on using this module with a ``nccl`` backend or a ``gloo`` backend (that uses Infiniband), together with a DataLoader that uses multiple workers, please change the multiprocessing start method to - ``forkserver`` (Python 3 only) or ``spawn``. Unfortunately + ``forkserver`` or ``spawn``. Unfortunately Gloo (that uses Infiniband) and NCCL2 are not fork safe, and you will likely experience deadlocks if you don't change this setting. @@ -1637,11 +1637,17 @@ def _get_active_ddp_module(cls): @contextmanager @torch._disable_dynamo(recursive=False) def _inside_ddp_forward(self): + # Save and restore the previous _active_ddp_module to handle nested + # DDP correctly (e.g., TorchRec wraps embeddings in an inner DDP inside + # an outer DDP). Without this, the inner DDP's exit would clear the + # flag to None, causing DDPOptimizer to miss compiled regions that run + # after the inner forward. + old = DistributedDataParallel._active_ddp_module DistributedDataParallel._active_ddp_module = self try: yield finally: - DistributedDataParallel._active_ddp_module = None + DistributedDataParallel._active_ddp_module = old def _run_ddp_forward(self, *inputs, **kwargs): if self._use_python_reducer: diff --git a/torch/nn/parallel/scatter_gather.py b/torch/nn/parallel/scatter_gather.py index 8954302f2a44e..bcd975a4fe2f7 100644 --- a/torch/nn/parallel/scatter_gather.py +++ b/torch/nn/parallel/scatter_gather.py @@ -58,19 +58,19 @@ def scatter_map(obj): if _is_namedtuple(obj): return [ type(obj)(*args) - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] for args in zip(*map(scatter_map, obj), strict=False) ] if isinstance(obj, tuple) and len(obj) > 0: - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] return list(zip(*map(scatter_map, obj), strict=False)) if isinstance(obj, list) and len(obj) > 0: - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] return [list(i) for i in zip(*map(scatter_map, obj), strict=False)] if isinstance(obj, dict) and len(obj) > 0: return [ type(obj)(i) - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] for i in zip(*map(scatter_map, obj.items()), strict=False) ] return [obj for _ in target_gpus] @@ -136,9 +136,9 @@ def gather_map(outputs): # pyrefly: ignore [not-callable] return type(out)((k, gather_map([d[k] for d in outputs])) for k in out) if _is_namedtuple(out): - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type] return type(out)._make(map(gather_map, zip(*outputs, strict=True))) - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type] return type(out)(map(gather_map, zip(*outputs, strict=True))) # Recursive function calls like this create reference cycles. diff --git a/torch/nn/parameter.py b/torch/nn/parameter.py index 824df5c8d317c..f5af0581d61fe 100644 --- a/torch/nn/parameter.py +++ b/torch/nn/parameter.py @@ -101,6 +101,7 @@ def __reduce_ex__(self, proto): (self.data, self.requires_grad, hooks, state), ) + # pyrefly: ignore [bad-override] __torch_function__ = _disabled_torch_function_impl @@ -273,6 +274,7 @@ def __new__(cls, data=None, *, persistent=True): t._is_buffer = True return t + # pyrefly: ignore [bad-override] __torch_function__ = _disabled_torch_function_impl diff --git a/torch/nn/utils/spectral_norm.py b/torch/nn/utils/spectral_norm.py index b785e080b79e6..79fea5a50b793 100644 --- a/torch/nn/utils/spectral_norm.py +++ b/torch/nn/utils/spectral_norm.py @@ -185,8 +185,6 @@ def apply( return fn -# This is a top level class because Py2 pickle doesn't like inner class nor an -# instancemethod. class SpectralNormLoadStateDictPreHook: # See docstring of SpectralNorm._version on the changes to spectral_norm. def __init__(self, fn) -> None: @@ -244,8 +242,6 @@ def __call__( state_dict[weight_key + "_v"] = v -# This is a top level class because Py2 pickle doesn't like inner class nor an -# instancemethod. class SpectralNormStateDictHook: # See docstring of SpectralNorm._version on the changes to spectral_norm. def __init__(self, fn) -> None: diff --git a/torch/onnx/_internal/exporter/_onnx_program.py b/torch/onnx/_internal/exporter/_onnx_program.py index c47ad25b31bf9..5d7484d85bed5 100644 --- a/torch/onnx/_internal/exporter/_onnx_program.py +++ b/torch/onnx/_internal/exporter/_onnx_program.py @@ -157,7 +157,7 @@ def _to_ort_value(input: torch.Tensor | int | float | str | bool) -> ort.OrtValu int: np.int64, float: np.float32, } - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] dtype = dtype_mapping.get(type(input)) return ort.OrtValue.ortvalue_from_numpy(np.array(input, dtype=dtype)) diff --git a/torch/onnx/_internal/torchscript_exporter/_type_utils.py b/torch/onnx/_internal/torchscript_exporter/_type_utils.py index acd13f914ae2d..c0ffa56a38e3b 100644 --- a/torch/onnx/_internal/torchscript_exporter/_type_utils.py +++ b/torch/onnx/_internal/torchscript_exporter/_type_utils.py @@ -173,6 +173,7 @@ def from_onnx_type( """ if onnx_type not in _ONNX_TO_SCALAR_TYPE: raise errors.OnnxExporterError(f"Unknown onnx_type: {onnx_type}") + # pyrefly: ignore [redundant-cast] return _ONNX_TO_SCALAR_TYPE[typing.cast(_C_onnx.TensorProtoDataType, onnx_type)] @classmethod diff --git a/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py b/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py index 2fcbdad947082..7beb3de13f272 100644 --- a/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py +++ b/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py @@ -167,7 +167,7 @@ def scaled_dot_product_attention( # NOTE: onnx-script has different logic here, because the attribute perms in # transpose needs list of ints key_shape_builtin = symbolic_helper._get_tensor_rank(key) - # pyrefly: ignore [no-matching-overload] + # pyrefly: ignore [bad-argument-type, no-matching-overload] key_transposed_axes = list(range(key_shape_builtin)) key_transposed_axes[-1], key_transposed_axes[-2] = ( key_transposed_axes[-2], diff --git a/torch/optim/_muon.py b/torch/optim/_muon.py index e441c8b911b2f..2e45e07c4a596 100644 --- a/torch/optim/_muon.py +++ b/torch/optim/_muon.py @@ -273,6 +273,29 @@ def step(self, closure=None): adjust_lr_fn (str, optional): function to adjust learning rate. One of "original" and "match_rms_adamw". If not specified, we will default to use "original". (default: None) + Example: + >>> # xdoctest: +SKIP + >>> # Muon only supports 2D params; use a standard optimizer + >>> # such as AdamW for biases, embeddings, and other non-2D + >>> # parameters. + >>> muon_params = [ + ... p for p in model.parameters() if p.ndim == 2 + ... ] + >>> other_params = [ + ... p for p in model.parameters() if p.ndim != 2 + ... ] + >>> optim_muon = torch.optim.Muon( + ... muon_params, lr=0.02, momentum=0.95 + ... ) + >>> optim_adamw = torch.optim.AdamW( + ... other_params, lr=3e-4, weight_decay=0.01 + ... ) + >>> optim_muon.zero_grad() + >>> optim_adamw.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optim_muon.step() + >>> optim_adamw.step() + .. _Muon\: An optimizer for hidden layers in neural networks: https://kellerjordan.github.io/posts/muon/ .. _Muon is Scalable for LLM Training: diff --git a/torch/optim/adagrad.py b/torch/optim/adagrad.py index a6a57fb61b8ba..d18e32fe17a98 100644 --- a/torch/optim/adagrad.py +++ b/torch/optim/adagrad.py @@ -16,6 +16,7 @@ _to_scalar, _use_grad_for_differentiable, _view_as_real, + DeviceDict, Optimizer, ParamsT, ) @@ -73,6 +74,7 @@ def __init__( if foreach: raise RuntimeError("`fused` and `foreach` cannot be `True` together.") self._need_device_dtype_check_for_fused = True + self._step_supports_amp_scaling = True for group in self.param_groups: for p in group["params"]: @@ -106,6 +108,20 @@ def __setstate__(self, state): group.setdefault("differentiable", False) fused = group.setdefault("fused", None) + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0 and not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = ( + torch.tensor( + step_val, + dtype=_get_scalar_dtype(is_fused=fused), + device=p.device, + ) + if group["fused"] + else torch.tensor(step_val, dtype=_get_scalar_dtype()) + ) + state_values = list(self.state.values()) step_is_tensor = (len(state_values) != 0) and torch.is_tensor( state_values[0]["step"] @@ -132,13 +148,38 @@ def _init_group(self, group, params_with_grad, grads, state_sums, state_steps): "_need_device_dtype_check_for_fused", True, ): - _device_dtype_check_for_fused(p, cuda_unsupported=True) + _device_dtype_check_for_fused(p) self._need_device_dtype_check_for_fused = False has_sparse_grad |= p.grad.is_sparse has_complex |= torch.is_complex(p) params_with_grad.append(p) grads.append(p.grad) state = self.state[p] + if len(state) == 0: + if group["fused"]: + _device_dtype_check_for_fused(p) + + state["step"] = ( + torch.zeros( + (), + dtype=_get_scalar_dtype(is_fused=group["fused"]), + device=p.device, + ) + if group["fused"] + else torch.tensor(0.0, dtype=_get_scalar_dtype()) + ) + + initial_accumulator_value = self.defaults[ + "initial_accumulator_value" + ] + init_value = ( + complex(initial_accumulator_value, initial_accumulator_value) + if torch.is_complex(p) + else initial_accumulator_value + ) + state["sum"] = torch.full_like( + p, init_value, memory_format=torch.preserve_format + ) state_sums.append(state["sum"]) state_steps.append(state["step"]) @@ -230,7 +271,7 @@ def step(self, closure=None): {_foreach_doc} {_maximize_doc} {_differentiable_doc} - fused (bool, optional): whether the fused implementation (CPU only) is used. + fused (bool, optional): whether the fused implementation (CPU and CUDA only) is used. Currently, `torch.float64`, `torch.float32`, `torch.float16`, and `torch.bfloat16` are supported. (default: None). Please note that the fused implementations does not support sparse or complex gradients. @@ -507,7 +548,7 @@ def _fused_adagrad( grad_scale: Tensor | None, found_inf: Tensor | None, *, - lr: float, + lr: float | Tensor, weight_decay: float, lr_decay: float, eps: float, @@ -526,12 +567,15 @@ def _fused_adagrad( "adagrad with fused=True does not support differentiable=True" ) - lr = _to_scalar(lr) - - grad_scale_dict = ( - {grad_scale.device: grad_scale} if grad_scale is not None else None + grad_scale_dict: DeviceDict = ( + {grad_scale.device: grad_scale} if grad_scale is not None else {} + ) + found_inf_dict: DeviceDict = ( + {found_inf.device: found_inf} if found_inf is not None else {} + ) + lr_dict: DeviceDict | None = ( + {lr.device: lr} if isinstance(lr, Tensor) and str(lr.device) != "cpu" else None ) - found_inf_dict = {found_inf.device: found_inf} if found_inf is not None else None grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( [params, grads, state_sums, state_steps] # type: ignore[list-item] @@ -551,14 +595,17 @@ def _fused_adagrad( device_state_steps = cast(list[Tensor], device_state_steps_) device_grad_scale, device_found_inf = None, None - if grad_scale is not None and grad_scale_dict is not None: - if device not in grad_scale_dict: - grad_scale_dict[device] = grad_scale.to(device, non_blocking=True) # type: ignore[index] - device_grad_scale = grad_scale_dict[device] # type: ignore[index] - if found_inf is not None and found_inf_dict is not None: - if found_inf not in found_inf_dict: - found_inf_dict[device] = found_inf.to(device, non_blocking=True) # type: ignore[index] - device_found_inf = found_inf_dict[device] # type: ignore[index] + if grad_scale is not None: + device_grad_scale = grad_scale_dict.setdefault( + device, grad_scale.to(device, non_blocking=True) + ) + if found_inf is not None: + device_found_inf = found_inf_dict.setdefault( + device, found_inf.to(device, non_blocking=True) + ) + if lr_dict is not None and device not in lr_dict: + lr_dict[device] = lr.to(device=device, non_blocking=True) # type: ignore[union-attr] + lr = lr_dict[device] torch._foreach_add_(device_state_steps, 1) torch._fused_adagrad_( device_params, diff --git a/torch/overrides.py b/torch/overrides.py index db945dd73322d..67920aef7d853 100644 --- a/torch/overrides.py +++ b/torch/overrides.py @@ -316,9 +316,6 @@ def get_ignored_functions() -> set[Callable]: torch.unify_type_list, torch.is_warn_always_enabled, torch.set_warn_always, - torch.vitals_enabled, - torch.set_vital, - torch.read_vitals, torch.vmap, torch.cond, torch.frombuffer, @@ -385,6 +382,8 @@ def get_ignored_functions() -> set[Callable]: Tensor._addmm_activation, Tensor.to_padded_tensor, Tensor._use_count, + Tensor._philox_normal_, + Tensor._philox_uniform_, } if sys.version_info >= (3, 14): diff --git a/torch/package/package_importer.py b/torch/package/package_importer.py index 1cb473e0aed68..a42af10794fa5 100644 --- a/torch/package/package_importer.py +++ b/torch/package/package_importer.py @@ -154,6 +154,7 @@ def __init__( # used for torch.serialization._load self.Unpickler = lambda *args, **kwargs: PackageUnpickler(self, *args, **kwargs) + # pyrefly: ignore [bad-override] def import_module(self, name: str, package=None): """Load a module from the package if it hasn't already been loaded, and then return the module. Modules are loaded locally diff --git a/torch/return_types.py b/torch/return_types.py index d456742be4b88..4c74fe01537e7 100644 --- a/torch/return_types.py +++ b/torch/return_types.py @@ -39,6 +39,7 @@ def structseq_unflatten(values, context): globals()[name] = _attr if not name.startswith("_"): + # pyrefly: ignore [unresolvable-dunder-all] __all__.append(name) all_return_types.append(_attr) diff --git a/torch/serialization.py b/torch/serialization.py index f6ebe0fce7e4d..32314ae0723e3 100644 --- a/torch/serialization.py +++ b/torch/serialization.py @@ -1384,7 +1384,7 @@ def load( second step is a no-op if the final location is CPU. When the ``mmap`` flag is set, instead of copying the tensor storages from disk to CPU memory in the first step, ``f`` is mapped, which means tensor storages will be lazily loaded when their data is accessed. - pickle_load_args: (Python 3 only) optional keyword arguments passed over to + pickle_load_args: optional keyword arguments passed over to :func:`pickle_module.load` and :func:`pickle_module.Unpickler`, only works if :attr:`weights_only=False`, e.g., :attr:`errors=...`. diff --git a/torch/sparse/__init__.py b/torch/sparse/__init__.py index 1e71d8fc4d5c2..78846265e9669 100644 --- a/torch/sparse/__init__.py +++ b/torch/sparse/__init__.py @@ -642,20 +642,27 @@ def convert_to_strided_representation(args): if obj.layout is torch.sparse_coo: # pyrefly: ignore [no-matching-overload] d.update( - indices=obj._indices(), is_coalesced=obj.is_coalesced() + # pyrefly: ignore [bad-argument-type] + indices=obj._indices(), + # pyrefly: ignore [bad-argument-type] + is_coalesced=obj.is_coalesced(), ) values = obj._values() elif obj.layout in {torch.sparse_csr, torch.sparse_bsr}: # pyrefly: ignore [no-matching-overload] d.update( + # pyrefly: ignore [bad-argument-type] compressed_indices=obj.crow_indices(), + # pyrefly: ignore [bad-argument-type] plain_indices=obj.col_indices(), ) values = obj.values() else: # pyrefly: ignore [no-matching-overload] d.update( + # pyrefly: ignore [bad-argument-type] compressed_indices=obj.ccol_indices(), + # pyrefly: ignore [bad-argument-type] plain_indices=obj.row_indices(), ) values = obj.values() diff --git a/torch/sparse/_semi_structured_ops.py b/torch/sparse/_semi_structured_ops.py index 6e12d76fc8abc..0254a33a4bc62 100644 --- a/torch/sparse/_semi_structured_ops.py +++ b/torch/sparse/_semi_structured_ops.py @@ -16,6 +16,8 @@ "semi_sparse_linear", "semi_sparse_scaled_mm", "semi_sparse_clone", + "semi_sparse_to", + "semi_sparse_to_copy", ] @@ -137,20 +139,14 @@ def semi_sparse_mm(func, types, args=(), kwargs=None) -> torch.Tensor: "`SparseSemiStructuredTensor` matmul: Broadcasting is not implemented" ) if isinstance(A, torch.sparse.SparseSemiStructuredTensor): - row, col = B.shape - B_padded = A._pad_dense_input(B) - res = A._mm(B_padded) - return res[:, :col] + return A._mm(B) else: B_t = B.t() if not isinstance(B_t, torch.sparse.SparseSemiStructuredTensor): raise AssertionError( f"expected SparseSemiStructuredTensor, got {type(B_t).__name__}" ) - row, col = A.shape - A_padded = B._pad_dense_input(A) - res = B_t._mm(A_padded.t()).t() - return res[:row, :] + return B_t._mm(A, should_transpose_dense=True).t() def semi_sparse_addmm(func, types, args=(), kwargs=None) -> torch.Tensor: @@ -175,9 +171,7 @@ def semi_sparse_addmm(func, types, args=(), kwargs=None) -> torch.Tensor: f"expected SparseSemiStructuredTensor, got {type(B_t).__name__}" ) row, _col = A.shape - A_padded = B_t._pad_dense_input(A) - result = B_t._mm(A_padded.t(), bias=bias).t() - return result[:row, :] + return B_t._mm(A, bias=bias, should_transpose_dense=True).t() def semi_sparse_linear(func, types, args=(), kwargs=None) -> torch.Tensor: @@ -188,7 +182,6 @@ def semi_sparse_linear(func, types, args=(), kwargs=None) -> torch.Tensor: shape = A.shape A_2d = A.view(-1, shape[-1]) - if bias is None: res = A_2d @ B.t() else: @@ -197,7 +190,6 @@ def semi_sparse_linear(func, types, args=(), kwargs=None) -> torch.Tensor: types=None, args=[bias, A_2d, B.t()], ) - return res.view(*shape[:-1], -1) @@ -263,3 +255,46 @@ def semi_sparse_clone(func, types, args=(), kwargs=None) -> torch.Tensor: alg_id_cusparselt=self.alg_id_cusparselt, requires_grad=self.requires_grad, ) + + +def semi_sparse_to_copy(func, types, args, kwargs=None) -> torch.Tensor: + self = args[0] + kwargs = kwargs or {} + + device = kwargs.get("device", None) + + if device is not None and torch.device(device).type == "cpu": + dense = self.to_dense() + return func(dense, **kwargs) + + raise NotImplementedError( + f"`_to_copy()` with kwargs={kwargs} is not implemented " + "for SparseSemiStructuredTensor. Only converting to CPU is supported currently." + ) + + +def semi_sparse_to(func, types, args, kwargs=None) -> torch.Tensor: + self = args[0] + remaining_args = args[1:] + kwargs = kwargs or {} + + # Determine the target device from args/kwargs + device = None + if remaining_args: + first_arg = remaining_args[0] + if isinstance(first_arg, (torch.device, str)): + try: + device = torch.device(first_arg) + except RuntimeError: + pass + if "device" in kwargs: + device = torch.device(kwargs["device"]) + + if device is not None and device.type == "cpu": + dense = self.to_dense() + return func(dense, *remaining_args, **kwargs) + + raise NotImplementedError( + f"`to()` with args={remaining_args}, kwargs={kwargs} is not implemented " + "for SparseSemiStructuredTensor. Only `to('cpu')` is supported currently." + ) diff --git a/torch/sparse/semi_structured.py b/torch/sparse/semi_structured.py index 7e9fb868c335f..4d99dac9eadef 100644 --- a/torch/sparse/semi_structured.py +++ b/torch/sparse/semi_structured.py @@ -19,6 +19,8 @@ semi_sparse_mm, semi_sparse_scaled_mm, semi_sparse_t, + semi_sparse_to, + semi_sparse_to_copy, semi_sparse_values, semi_sparse_view, ) @@ -231,9 +233,10 @@ def _load_dispatch_table(cls, custom_dispatch_table=None) -> None: torch.ops.aten.matmul: semi_sparse_mm, torch.ops.aten.addmm: semi_sparse_addmm, torch.ops.aten.linear: semi_sparse_linear, - torch.ops.aten._to_copy: fallback_dispatcher, + torch.ops.aten._to_copy: semi_sparse_to_copy, torch.ops.aten._scaled_mm: semi_sparse_scaled_mm, torch.ops.aten.clone: semi_sparse_clone, + torch.ops.aten.to: semi_sparse_to, } if custom_dispatch_table is not None: cls.SPARSE_DISPATCH.update(custom_dispatch_table) @@ -281,35 +284,16 @@ def _validate_device_dim_dtype_shape(cls, original_tensor: torch.Tensor) -> None f"Both dimensions must be larger or equal than and a multiple of ({min_rows}, {min_cols})" ) - @classmethod - def _pad_dense_input(cls, dense_input: torch.Tensor) -> torch.Tensor: - """ - Calculates padding for dense tensor and pads tensor if necessary. - If padding is not required, this function returns the original tensor. - """ - # only 2d matmul - if dense_input.dim() != 2: - raise AssertionError(f"dense_input must be 2D, got {dense_input.dim()}D") - - # check shape - m, n = dense_input.shape - min_rows = cls._DTYPE_SHAPE_CONSTRAINTS[dense_input.dtype].dense_min_rows - min_cols = cls._DTYPE_SHAPE_CONSTRAINTS[dense_input.dtype].dense_min_cols - - # calculate padding - to_pad_m = -m % min_rows if m < min_rows or m % min_rows else 0 - to_pad_n = -n % min_cols if n < min_cols or n % min_rows else 0 - if to_pad_m or to_pad_n: - return torch.nn.functional.pad(dense_input, (0, to_pad_n, 0, to_pad_m)) - else: - return dense_input - def to_dense(self): # type:ignore[override] col = self.shape[-1] return torch.mm(self, torch.eye(col, dtype=self.dtype, device=self.device)) @classmethod - def from_dense(cls, original_tensor: torch.Tensor) -> "SparseSemiStructuredTensor": + def from_dense( + cls, + original_tensor: torch.Tensor, + alg_id: int = _DEFAULT_ALG_ID, + ) -> "SparseSemiStructuredTensor": raise NotImplementedError def _mm( @@ -325,6 +309,7 @@ def _mm( def to_sparse_semi_structured( original_tensor: torch.Tensor, transposed: bool = False, + alg_id: int = SparseSemiStructuredTensor._DEFAULT_ALG_ID, ) -> SparseSemiStructuredTensor: """ This function converts a dense tensor into a sparse semi-structured tensor. @@ -338,6 +323,8 @@ def to_sparse_semi_structured( Args: original_tensor (Tensor): the dense tensor to convert transposed (bool, optional): deprecated arg to be removed in another release. Do not use. + alg_id (int, optional): the algorithm id to use for cuSPARSELt matmul. Defaults to 0. + Can be obtained via ``torch._cslt_sparse_mm_search``. Returns: SparseSemiStructuredTensor: A sparse semi-structured tensor created from the given original_tensor Raises: @@ -387,7 +374,7 @@ def to_sparse_semi_structured( else torch.sparse.SparseSemiStructuredTensorCUSPARSELT ) - return SPARSE_SUBCLASS.from_dense(original_tensor) + return SPARSE_SUBCLASS.from_dense(original_tensor, alg_id=alg_id) class SparseSemiStructuredTensorCUTLASS(SparseSemiStructuredTensor): @@ -412,7 +399,9 @@ class SparseSemiStructuredTensorCUTLASS(SparseSemiStructuredTensor): @classmethod def from_dense( - cls, original_tensor: torch.Tensor + cls, + original_tensor: torch.Tensor, + alg_id: int = SparseSemiStructuredTensor._DEFAULT_ALG_ID, ) -> "SparseSemiStructuredTensorCUTLASS": cls._validate_device_dim_dtype_shape(original_tensor) ( @@ -519,7 +508,12 @@ def prune_dense_static_sort( ) def _mm( - self, B: torch.Tensor, *, bias: torch.Tensor | None = None, **kwargs + self, + B: torch.Tensor, + *, + bias: torch.Tensor | None = None, + should_transpose_dense: bool = False, + **kwargs, ) -> torch.Tensor: if isinstance(B, SparseSemiStructuredTensor): raise ValueError( @@ -535,13 +529,18 @@ def _mm( f"`{cls_name}` matmul: operation is not supported" ) else: - if bias is None: - res = torch._sparse_semi_structured_mm(self.packed, self.meta, B) - else: - res = torch._sparse_semi_structured_addmm( - bias, self.packed, self.meta, B - ) - return res[: self.shape[0]] + _ensure_cutlass_mm_registered() + constraints = self._DTYPE_SHAPE_CONSTRAINTS[B.dtype] + return torch.ops.semi_structured.cutlass_mm( + B, + self.packed, + self.meta, + bias, + self.shape[0], + constraints.dense_min_rows, + constraints.dense_min_cols, + should_transpose_dense, + ) class SparseSemiStructuredTensorCUSPARSELT(SparseSemiStructuredTensor): @@ -566,7 +565,9 @@ class SparseSemiStructuredTensorCUSPARSELT(SparseSemiStructuredTensor): @classmethod def from_dense( - cls, original_tensor: torch.Tensor + cls, + original_tensor: torch.Tensor, + alg_id: int = SparseSemiStructuredTensor._DEFAULT_ALG_ID, ) -> "SparseSemiStructuredTensorCUSPARSELT": cls._validate_device_dim_dtype_shape(original_tensor) # pyrefly: ignore [no-matching-overload] @@ -578,7 +579,7 @@ def from_dense( meta_t=None, compressed_swizzled_bitmask=None, fuse_transpose_cusparselt=SparseSemiStructuredTensor._FUSE_TRANSPOSE, - alg_id_cusparselt=SparseSemiStructuredTensor._DEFAULT_ALG_ID, + alg_id_cusparselt=alg_id, requires_grad=original_tensor.requires_grad, ) @@ -648,7 +649,12 @@ def prune_dense_static_sort( ) def _mm( - self, B: torch.Tensor, *, bias: torch.Tensor | None = None, **kwargs + self, + B: torch.Tensor, + *, + bias: torch.Tensor | None = None, + should_transpose_dense: bool = False, + **kwargs, ) -> torch.Tensor: if isinstance(B, SparseSemiStructuredTensor): raise ValueError( @@ -682,11 +688,153 @@ def _mm( f"`{self.__class__.__name__}` matmul: operation is not supported" ) else: - res = torch._cslt_sparse_mm( - self.packed, + _ensure_cusparselt_mm_registered() + constraints = self._DTYPE_SHAPE_CONSTRAINTS[B.dtype] + return torch.ops.semi_structured.cusparselt_mm( B, - bias=bias, - transpose_result=self.fuse_transpose_cusparselt, - alg_id=self.alg_id_cusparselt, + self.packed, + bias, + self.shape[0], + constraints.dense_min_rows, + constraints.dense_min_cols, + self.fuse_transpose_cusparselt, + self.alg_id_cusparselt, + should_transpose_dense, ) - return res.t() if self.fuse_transpose_cusparselt else res + + +_cutlass_mm_registered = False + + +def _ensure_cutlass_mm_registered(): + """Lazily register the cutlass_mm custom op. + + Registration is deferred to avoid importing torch.library at module load + time, since torch.sparse is imported early during ``import torch``. + """ + global _cutlass_mm_registered + if _cutlass_mm_registered: + return + _cutlass_mm_registered = True + + from torch.library import custom_op + + @custom_op("semi_structured::cutlass_mm", mutates_args=()) + def cutlass_mm( + dense: torch.Tensor, + packed: torch.Tensor, + meta: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + min_rows: int, + min_cols: int, + should_transpose_dense: bool, + ) -> torch.Tensor: + m, n = dense.shape + to_pad_m = (-m) % min_rows + to_pad_n = (-n) % min_cols + need_pad = to_pad_m != 0 or to_pad_n != 0 + dense_padded = dense + if need_pad: + dense_padded = torch.nn.functional.pad(dense, (0, to_pad_n, 0, to_pad_m)) + mm_input = dense_padded.t() if should_transpose_dense else dense_padded + if bias is None: + res = torch._sparse_semi_structured_mm(packed, meta, mm_input) + else: + res = torch._sparse_semi_structured_addmm(bias, packed, meta, mm_input) + if need_pad: + out_cols = m if should_transpose_dense else n + return ( + res[:out_features] + .narrow(1, 0, out_cols) + .clone(memory_format=torch.contiguous_format) + ) + return res.contiguous() + + @cutlass_mm.register_fake + def _cutlass_mm_fake( + dense: torch.Tensor, + packed: torch.Tensor, + meta: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + min_rows: int, + min_cols: int, + transpose_dense: bool, + ) -> torch.Tensor: + out_cols = dense.shape[0] if transpose_dense else dense.shape[1] + return torch.empty( + out_features, + out_cols, + dtype=dense.dtype, + device=dense.device, + ) + + +_cusparselt_mm_registered = False + + +def _ensure_cusparselt_mm_registered(): + """Lazily register the cusparselt_mm custom op.""" + global _cusparselt_mm_registered + if _cusparselt_mm_registered: + return + _cusparselt_mm_registered = True + + from torch.library import custom_op + + @custom_op("semi_structured::cusparselt_mm", mutates_args=()) + def cusparselt_mm( + dense: torch.Tensor, + packed: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + min_rows: int, + min_cols: int, + fuse_transpose: bool, + alg_id: int, + should_transpose_dense: bool = False, + ) -> torch.Tensor: + m, n = dense.shape + to_pad_m = (-m) % min_rows + to_pad_n = (-n) % min_cols + need_pad = to_pad_m != 0 or to_pad_n != 0 + dense_padded = dense + if need_pad: + dense_padded = torch.nn.functional.pad(dense, (0, to_pad_n, 0, to_pad_m)) + mm_input = dense_padded.t() if should_transpose_dense else dense_padded + res = torch._cslt_sparse_mm( + packed, + mm_input, + bias=bias, + transpose_result=fuse_transpose, + alg_id=alg_id, + ) + if fuse_transpose: + res = res.t() + if need_pad: + out_cols = m if should_transpose_dense else n + return res.narrow(1, 0, out_cols).clone( + memory_format=torch.contiguous_format + ) + return res.contiguous() + + @cusparselt_mm.register_fake + def _cusparselt_mm_fake( + dense: torch.Tensor, + packed: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + min_rows: int, + min_cols: int, + fuse_transpose: bool, + alg_id: int, + should_transpose_dense: bool, + ) -> torch.Tensor: + out_cols = dense.shape[0] if should_transpose_dense else dense.shape[1] + return torch.empty( + out_features, + out_cols, + dtype=dense.dtype, + device=dense.device, + ) diff --git a/torch/testing/_comparison.py b/torch/testing/_comparison.py index 2605411ed7c9d..ef60dd19f7c3e 100644 --- a/torch/testing/_comparison.py +++ b/torch/testing/_comparison.py @@ -18,6 +18,37 @@ HAS_NUMPY = False np = None # type: ignore[assignment] +_HAS_DTENSOR = torch.distributed.is_available() + + +def _unwrap_dtensor_for_comparison(actual, expected): + """Handle DTensor inputs for assertEqual/assert_close.""" + if not _HAS_DTENSOR: + return actual, expected + from torch.distributed.tensor import DTensor + + actual_dt = isinstance(actual, DTensor) + expected_dt = isinstance(expected, DTensor) + if actual_dt and expected_dt: + if actual.placements != expected.placements: + raise AssertionError( + f"DTensor placements do not match: " + f"{actual.placements} != {expected.placements}" + ) + if actual.device_mesh != expected.device_mesh: + raise AssertionError( + f"DTensor device meshes do not match: " + f"{actual.device_mesh} != {expected.device_mesh}" + ) + return actual.to_local(), expected.to_local() + elif actual_dt != expected_dt: + raise TypeError( + "Comparing a DTensor to a non-DTensor is ambiguous. " + "Call .full_tensor() to compare the full logical tensor " + "or .to_local() to compare the local shard." + ) + return actual, expected + class ErrorMeta(Exception): """Internal testing exception that makes that carries error metadata.""" @@ -1573,6 +1604,8 @@ def assert_close( # Hide this function from `pytest`'s traceback __tracebackhide__ = True + actual, expected = _unwrap_dtensor_for_comparison(actual, expected) + error_metas = not_close_error_metas( actual, expected, diff --git a/torch/testing/_internal/common_cuda.py b/torch/testing/_internal/common_cuda.py index 0dc3d0b4f10cb..8c2f0fc64fef7 100644 --- a/torch/testing/_internal/common_cuda.py +++ b/torch/testing/_internal/common_cuda.py @@ -5,7 +5,7 @@ import functools import torch import torch.cuda -from torch.testing._internal.common_utils import LazyVal, TEST_NUMBA, TEST_WITH_ROCM, TEST_CUDA, IS_WINDOWS, IS_MACOS +from torch.testing._internal.common_utils import LazyVal, TEST_NUMBA, TEST_WITH_ROCM, TEST_CUDA, IS_WINDOWS, IS_MACOS, TEST_XPU import inspect import contextlib import os @@ -43,6 +43,7 @@ IS_SM89 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (8, 9)) IS_SM90 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0)) IS_SM100 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0)) +IS_SM12X = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 12) @contextlib.contextmanager def blas_library_context(backend): @@ -76,8 +77,16 @@ def evaluate_platform_supports_flash_attention(): return evaluate_gfx_arch_within(arch_list) if TEST_CUDA: return not IS_WINDOWS and SM80OrLater + if TEST_XPU: + return True return False +def evaluate_platform_supports_ck_sdpa(): + if TEST_WITH_ROCM: + return torch.backends.cuda.is_ck_sdpa_available() + else: + return False + def evaluate_platform_supports_efficient_attention(): if TEST_WITH_ROCM: arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950"] @@ -86,6 +95,8 @@ def evaluate_platform_supports_efficient_attention(): return evaluate_gfx_arch_within(arch_list) if TEST_CUDA: return True + if TEST_XPU: + return True return False def evaluate_platform_supports_cudnn_attention(): @@ -111,12 +122,16 @@ def evaluate_platform_supports_green_context(): PLATFORM_SUPPORTS_FUSED_SDPA: bool = TEST_CUDA and not TEST_WITH_ROCM +PLATFORM_SUPPORTS_CK_SDPA: bool = LazyVal(lambda: evaluate_platform_supports_ck_sdpa()) + def evaluate_platform_supports_bf16(): if torch.version.cuda: return SM80OrLater elif torch.version.hip: return True + elif TEST_XPU: + return True return False @@ -140,6 +155,18 @@ def evaluate_platform_supports_half_atomics(): PLATFORM_SUPPORTS_GREEN_CONTEXT: bool = LazyVal(lambda: evaluate_platform_supports_green_context()) +def evaluate_platform_supports_workqueue_config(): + if IS_WINDOWS: + return False + if not _get_torch_cuda_version() >= (13, 1): + return False + driver_version = torch.utils.collect_env.get_nvidia_driver_version(torch.utils.collect_env.run) + if driver_version is None: + return False + return int(driver_version.split('.')[0]) >= 590 + +PLATFORM_SUPPORTS_WORKQUEUE_CONFIG: bool = LazyVal(lambda: evaluate_platform_supports_workqueue_config()) + def evaluate_platform_supports_fp8(): if torch.cuda.is_available(): if torch.version.hip: @@ -153,7 +180,10 @@ def evaluate_platform_supports_fp8(): return True else: return SM90OrLater or torch.cuda.get_device_capability() == (8, 9) - return False + if torch.xpu.is_available(): + return True + # As CPU supports FP8 and is always available, return True. + return True def evaluate_platform_supports_fp8_grouped_gemm(): if torch.cuda.is_available(): @@ -183,8 +213,21 @@ def evaluate_platform_supports_mxfp8_grouped_gemm(): return built_with_mslk and IS_SM100 return False +def evaluate_platform_supports_fp8_sparse(): + if torch.cuda.is_available(): + if torch.version.hip: + return 'gfx950' in torch.cuda.get_device_properties(0).gcnArchName + else: + return ( + (SM90OrLater or torch.cuda.get_device_capability() == (8, 9)) + and torch.backends.cusparselt.is_available() + and torch.backends.cusparselt.version() >= 602 + ) + return False + PLATFORM_SUPPORTS_MX_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_mx_gemm()) PLATFORM_SUPPORTS_FP8: bool = LazyVal(lambda: evaluate_platform_supports_fp8()) +PLATFORM_SUPPORTS_FP8_SPARSE: bool = LazyVal(lambda: evaluate_platform_supports_fp8_sparse()) PLATFORM_SUPPORTS_FP8_GROUPED_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_fp8_grouped_gemm()) PLATFORM_SUPPORTS_MXFP8_GROUPED_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_mxfp8_grouped_gemm()) @@ -192,7 +235,7 @@ def evaluate_platform_supports_mxfp8_grouped_gemm(): try: import numba.cuda TEST_NUMBA_CUDA = numba.cuda.is_available() - except (ImportError, RuntimeError): + except (ImportError, RuntimeError, OSError): TEST_NUMBA_CUDA = False TEST_NUMBA = False else: @@ -434,6 +477,9 @@ def xfailIfSM100OrLater(func): def xfailIfSM120OrLater(func): return func if not SM120OrLater else unittest.expectedFailure(func) +def xfailIfSM12X(func): + return func if not IS_SM12X else unittest.expectedFailure(func) + def xfailIfDistributedNotSupported(func): return func if not (IS_MACOS or IS_JETSON) else unittest.expectedFailure(func) diff --git a/torch/testing/_internal/common_device_type.py b/torch/testing/_internal/common_device_type.py index e2762861ccdac..8fc9f2cd9682f 100644 --- a/torch/testing/_internal/common_device_type.py +++ b/torch/testing/_internal/common_device_type.py @@ -17,6 +17,7 @@ import torch from torch._inductor.utils import GPU_TYPES +from torch._utils import _is_privateuse1_backend_available from torch.testing._internal.common_cuda import ( _get_torch_cuda_version, _get_torch_hipblaslt_version, @@ -33,7 +34,6 @@ get_tracked_input, IS_FBCODE, IS_MACOS, - is_privateuse1_backend_available, IS_REMOTE_GPU, IS_S390X, IS_SANDCASTLE, @@ -318,6 +318,15 @@ def _update_param_kwargs(param_kwargs, name, value): class DeviceTypeTestBase(TestCase): device_type: str = "generic_device_type" + # When True, @onlyOn-based decorators (@onlyCUDA, @onlyMPS, etc.) will not + # skip tests for this device type. This is a pragmatic short-term solution to + # allow PrivateUse1 backends to run tests that are currently gated behind + # device-specific decorators. It is intended to be used together with the + # skip mechanism (see https://github.com/pytorch/pytorch/issues/177253). + # In the longer term, we are incrementally migrating accelerator tests to be + # device-generic and removing @onlyCUDA on tests that should be device-generic. + bypass_device_restrictions: bool = False + # Flag to disable test suite early due to unrecoverable error such as CUDA error. _stop_test_suite = False @@ -670,6 +679,7 @@ class PrivateUse1TestBase(DeviceTypeTestBase): primary_device: ClassVar[str] device_mod = None device_type = "privateuse1" + bypass_device_restrictions = False @classmethod def get_primary_device(cls): @@ -723,7 +733,7 @@ def get_device_type_test_bases(): if torch.cuda.is_available(): test_bases.append(CUDATestBase) - if is_privateuse1_backend_available(): + if _is_privateuse1_backend_available(): test_bases.append(PrivateUse1TestBase) # Disable MPS testing in generic device testing temporarily while we're # ramping up support. @@ -750,7 +760,7 @@ def filter_desired_device_types(device_type_test_bases, except_for=None, only_fo # This handles the case where PrivateUse1TestBase.device_type has been # changed from "privateuse1" to the actual backend name (e.g., "openreg") # by setUpClass being called during previous instantiate_device_type_tests calls - if is_privateuse1_backend_available(): + if _is_privateuse1_backend_available(): privateuse1_backend_name = torch._C._get_privateuse1_backend_name() def func_replace(x: str) -> str: @@ -1444,6 +1454,8 @@ def __call__(self, fn): @wraps(fn) def only_fn(slf, *args, **kwargs): if slf.device_type not in self.device_type: + if getattr(slf, "bypass_device_restrictions", False): + return fn(slf, *args, **kwargs) reason = f"Only runs on {self.device_type}" if IS_SANDCASTLE or IS_FBCODE: print( diff --git a/torch/testing/_internal/common_distributed.py b/torch/testing/_internal/common_distributed.py index 86e36e5b094d0..894acb6be19fe 100644 --- a/torch/testing/_internal/common_distributed.py +++ b/torch/testing/_internal/common_distributed.py @@ -609,7 +609,10 @@ def create_tcp_store( TIMEOUT_DEFAULT = 500 else: TIMEOUT_DEFAULT = int(os.getenv("DISTRIBUTED_TESTS_DEFAULT_TIMEOUT", "300")) -TIMEOUT_OVERRIDE = {"test_ddp_uneven_inputs": 400} +TIMEOUT_OVERRIDE = { + "test_ddp_uneven_inputs": 400, + "test_DistributedDataParallel": 500, +} # https://github.com/pytorch/pytorch/issues/75665 @@ -737,6 +740,32 @@ def cleanup_temp_dir() -> None: tmp_dir.cleanup() +def retrieve_result_from_completion_queue( + process: torch.multiprocessing.Process, + completion_queue: torch.multiprocessing.Queue, + timeout: int | None = None, +) -> Any: + """Get result from the completion_queue associated with process. + + When the process finished without putting a result or the timeout expired an exception instance will be returned""" + queue_timeout = 120 if timeout is None else max(10, min(120, timeout // 4)) + start_time = time.time() + # Periodically check the process for liveness + while True: + try: + return completion_queue.get(timeout=queue_timeout) + except queue.Empty: + # If the process is no longer alive we cannot get a result from the queue unless it is there right now. + # This can happen if the timeout occurred just before the process put its result and terminated. + # So do a last check for emptiness before considering it as a failure. + if not process.is_alive() and completion_queue.empty(): + return RuntimeError(f"Exited with {process.exitcode}") + if timeout is not None: + elapsed = time.time() - start_time + if elapsed > timeout: + return RuntimeError(f"Process timed out out after {elapsed}s") + + # Most tests operate with this worldsize if TEST_WITH_ROCM: DEFAULT_WORLD_SIZE = min(4, max(2, torch.cuda.device_count())) @@ -945,7 +974,7 @@ def run_test(self, test_name: str, parent_pipe) -> None: "Process %s skipping test %s for following reason: %s", self.rank, test_name, - str(se), + se, ) sys.exit(TEST_SKIPS["generic"].exit_code) except Exception: @@ -1496,7 +1525,7 @@ def _check_return_codes(cls, failed_ranks, timeout, fn): "Thread %s skipping test %s for following reason: %s", rank, fn, - str(exc), + exc, ) if skip_code < 0: skip_code = TEST_SKIPS["generic"].exit_code @@ -2019,8 +2048,12 @@ def wrapper(self): # Drain all completion queues before raising any exception, # so stale results don't desync subsequent tests. deferred_exception = None - for i, completion_queue in enumerate(self.completion_queues): - rv = completion_queue.get() + for i, (p, completion_queue) in enumerate( + zip(self.processes, self.completion_queues) + ): + rv = retrieve_result_from_completion_queue( + p, completion_queue, timeout=get_timeout(self.id()) + ) if deferred_exception is not None: # Already captured an exception; just drain continue diff --git a/torch/testing/_internal/common_dtype.py b/torch/testing/_internal/common_dtype.py index c5f3ad3f390dc..6e62009b6b601 100644 --- a/torch/testing/_internal/common_dtype.py +++ b/torch/testing/_internal/common_dtype.py @@ -229,6 +229,13 @@ def highest_precision_float(device): return torch.float64 +def highest_precision_complex(device): + if torch.device(device).type == "mps": + return torch.complex64 + else: + return torch.complex128 + + float_to_corresponding_complex_type_map = { torch.float16: torch.complex32, torch.float32: torch.complex64, diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index db6904ecc0db3..9b1665c67d0b3 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -23,6 +23,7 @@ _dispatch_dtypes, floating_types, floating_types_and, complex_types, floating_and_complex_types, floating_and_complex_types_and, all_types_and_complex_and, all_types_and, all_types_and_complex, integral_types_and, empty_types, complex_types_and, integral_types, custom_types, all_types_complex_float8_and, float8_types, + highest_precision_complex, highest_precision_float, ) from torch.testing._internal.common_device_type import ( @@ -1616,7 +1617,7 @@ def sample_inputs_like_fns(self, device, dtype, requires_grad, **kwargs): ((S,), {'dtype': dtype, 'device': device}), # Hard-code some dtypes/devices. We want to test cases where the # (dtype, device) is different from the input's (dtype, device) - ((S,), {'dtype': torch.double if device != 'mps:0' else torch.float}), + ((S,), {'dtype': highest_precision_float(device)}), ((S,), {'device': 'cpu'}), ((S,), {'dtype': torch.double, 'device': 'cpu'}), ] @@ -1810,7 +1811,7 @@ def sample_inputs_new_fns(self, device, dtype, requires_grad, *, is_strided=Fals ((S,), (2, 3), (7, 8), {'dtype': dtype, 'device': device}), # Hard-code some dtypes/devices. We want to test cases where the # (dtype, device) is different from the input's (dtype, device) - ((S,), (10,), (S,), {'dtype': torch.double if device != 'mps:0' else torch.float}), + ((S,), (10,), (S,), {'dtype': highest_precision_float(device)}), ((S,), (1, 1, 12), (S, L, M), {'device': 'cpu'}), ((S,), (2, 2, 2), (L, M, S), {'dtype': torch.double, 'device': 'cpu'}), ] @@ -1933,7 +1934,7 @@ def sample_inputs_full_like(self, device, dtype, requires_grad, **kwargs): def get_val(dtype): return make_tensor([], dtype=dtype, device="cpu").item() - double_dtype = torch.double if device != "mps:0" else torch.float + double_dtype = highest_precision_float(device) inputs = [ ((), get_val(dtype), {}), ((S, S), get_val(dtype), {}), @@ -4058,8 +4059,8 @@ def sample_inputs_conv1d(op_info, device, dtype, requires_grad, **kwargs): def error_inputs_conv1d(opinfo, device, **kwargs): - dtype = torch.float64 if device != 'mps:0' else torch.float32 - cdtype = torch.complex128 if device != 'mps:0' else torch.complex64 + dtype = highest_precision_float(device) + cdtype = highest_precision_complex(device) make_arg = partial(make_tensor, device=device, dtype=dtype) make_int_arg = partial(make_tensor, device=device, dtype=torch.int64) make_complex_arg = partial(make_tensor, device=device, dtype=cdtype) @@ -4120,8 +4121,8 @@ def error_inputs_conv1d(opinfo, device, **kwargs): def error_inputs_conv2d(opinfo, device, **kwargs): - dtype = torch.float64 if device != 'mps:0' else torch.float32 - cdtype = torch.complex128 if device != 'mps:0' else torch.complex64 + dtype = highest_precision_float(device) + cdtype = highest_precision_complex(device) make_arg = partial(make_tensor, device=device, dtype=dtype) make_int_arg = partial(make_tensor, device=device, dtype=torch.int64) make_complex_arg = partial(make_tensor, device=device, dtype=cdtype) @@ -4259,8 +4260,8 @@ def sample_inputs_conv3d(opinfo, device, dtype, requires_grad, **kwargs): def error_inputs_conv3d(opinfo, device, **kwargs): - dtype = torch.float64 if device != 'mps:0' else torch.float32 - cdtype = torch.complex128 if device != 'mps:0' else torch.complex64 + dtype = highest_precision_float(device) + cdtype = highest_precision_complex(device) make_arg = partial(make_tensor, device=device, dtype=dtype) make_int_arg = partial(make_tensor, device=device, dtype=torch.int64) make_complex_arg = partial(make_tensor, device=device, dtype=cdtype) @@ -7000,8 +7001,6 @@ def skips_mvlgamma(skip_redundant=False): DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_small', dtypes=(torch.int8,)), - # NotImplementedError: The operator 'aten::mvlgamma.out' is not currently implemented for the MPS device - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), ) if skip_redundant: # Redundant tests @@ -9780,6 +9779,8 @@ def _should_disable_fastpath(self, opinfo, rightmost_arg, rightmost_arg_type, dt if self.arity == 1: if "foreach_abs" in opinfo.name and dtype in complex_types(): return True + if "foreach_clone" in opinfo.name: + return False # unary if opinfo.ref in (torch.abs, torch.neg): return False @@ -9958,7 +9959,7 @@ def sample_zero_size_tensor_inputs(self, opinfo, device, dtype, requires_grad, * for ord in (0, 1, 2, -1, -2, float('inf'), float('-inf')): input = sample_inputs_foreach(None, device, dtype, NUM_SIZE0_TENSORS, zero_size=True, **_foreach_inputs_kwargs) disable_fastpath = True - if ord in (1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16): + if ord in (0, 1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16): disable_fastpath = False yield ForeachSampleInput(input, ord=ord, disable_fastpath=disable_fastpath) @@ -9973,7 +9974,7 @@ def __call__(self, opinfo, device, dtype, requires_grad, **kwargs): for num_tensors, ord, out_dtype, intersperse_empty_tensors in product( num_input_tensors, (0, 1, 2, -1, -2, float('inf'), float('-inf')), - (None,) + (torch.complex128,) if dtype in complex_types() else (torch.float64,), + (None,) + (highest_precision_complex(device),) if dtype in complex_types() else (highest_precision_float(device),), (True, False), ): # inf norm and negative norms on empty tensors is not supported by our reference func vector norm: @@ -9984,7 +9985,7 @@ def __call__(self, opinfo, device, dtype, requires_grad, **kwargs): _foreach_inputs_kwargs["intersperse_empty_tensors"] = intersperse_empty_tensors input = sample_inputs_foreach(None, device, dtype, num_tensors, zero_size=False, **_foreach_inputs_kwargs) disable_fastpath = True - if ord in (1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16): + if ord in (0, 1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16): disable_fastpath = False yield ForeachSampleInput(input, ord=ord, disable_fastpath=disable_fastpath, dtype=out_dtype) @@ -10001,7 +10002,7 @@ def __call__(self, opinfo, device, dtype, requires_grad, **kwargs): for input in nan_inputs: x = torch.tensor(input, device=device) disable_fastpath = True - if ord in (1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16): + if ord in (0, 1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16): disable_fastpath = False yield ForeachSampleInput([x], ord=ord, disable_fastpath=disable_fastpath) @@ -11247,6 +11248,13 @@ def __call__(self, opinfo, device, dtype, requires_grad, **kwargs): ), ), ), + ForeachFuncInfo( + "clone", + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + supports_forward_ad=True, + supports_autograd=True, + supports_inplace_autograd=True, + ), ] foreach_binary_op_db: list[OpInfo] = [ @@ -13048,6 +13056,9 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', device_type='cuda', dtypes=[torch.cfloat], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), )), @@ -13549,6 +13560,9 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', 'test_reference_numerics_large', device_type='cuda', dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), )), UnaryUfuncInfo('cosh', ref=np_unary_ufunc_integer_promotion_wrapper(np.cosh), @@ -13581,6 +13595,9 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', 'test_reference_numerics_large', device_type='cuda', dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), )), OpInfo('cov', dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), @@ -13631,21 +13648,8 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k skips=( # cumsum does not handle correctly out= dtypes DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), - # The following dtypes did not work in forward but are listed by the OpInfo: {torch.complex64} + # The following dtypes did not work in forward but are listed by the OpInfo: {torch.bool} DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), - # RuntimeError: cumulative ops are not yet supported for complex - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', - device_type='mps', dtypes=(torch.complex64,) - ), - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_out_requires_grad_error', - device_type='mps', dtypes=(torch.complex64,) - ), - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', - device_type='mps', dtypes=(torch.complex64,) - ), ), sample_inputs_func=sample_inputs_cumulative_ops), OpInfo('cumprod', @@ -13656,21 +13660,8 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k skips=( # cumprod does not handle correctly out= dtypes DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), - # The following dtypes did not work in forward but are listed by the OpInfo: {torch.complex64} + # The following dtypes did not work in forward but are listed by the OpInfo: {torch.bool} DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), - # RuntimeError: cumulative ops are not yet supported for complex - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', - device_type='mps', dtypes=(torch.complex64,) - ), - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_out_requires_grad_error', - device_type='mps', dtypes=(torch.complex64,) - ), - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', - device_type='mps', dtypes=(torch.complex64,) - ), ), # gradgradcheck fails in fast_mode=True: #56275 sample_inputs_func=sample_inputs_cumprod, @@ -13822,6 +13813,8 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k # Reference: https://github.com/pytorch/pytorch/issues/48010 DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS), ), assert_autodiffed=True, supports_forward_ad=True, @@ -15425,7 +15418,9 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), decorators=(onlyNativeDeviceTypes,), - supports_autograd=False, + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, sample_inputs_func=sample_inputs_aminmax, skips=( # Exception: MPS supports tensors with dimensions <= 16, but got 65. @@ -16263,7 +16258,7 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k aliases=('group_norm',), ref=reference_group_norm, dtypes=floating_types_and(torch.float16, torch.bfloat16), - dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16, torch.int32, torch.int16), + dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16, torch.int32, torch.int16, torch.int8, torch.uint8), supports_out=False, supports_forward_ad=True, supports_fwgrad_bwgrad=True, @@ -16278,6 +16273,9 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k "test_comprehensive", device_type="cpu" ), + # MPS supports int8/uint8 but CPU does not, so consistency test cannot run + DecorateInfo(unittest.expectedFailure, 'TestConsistency', 'test_output_match', + device_type='mps', dtypes=(torch.int8, torch.uint8)), ], sample_inputs_func=sample_inputs_group_norm, reference_inputs_func=reference_inputs_group_norm, @@ -16632,7 +16630,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k supports_out=False, skips=( DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), - DecorateInfo(unittest.expectedFailure, 'TestDTensorOps', 'test_dtensor_op_db'), DecorateInfo(unittest.expectedFailure, 'TestInductorOpInfo', 'test_comprehensive'), DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), # Error: The operator 'aten::_upsample_bilinear2d_aa_backward.grad_input' @@ -17059,9 +17056,8 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k supports_expanded_weight=True, skips=( # RuntimeError: MPS device does not support linear for non-float inputs - # RuntimeError: mps linear does not support complex types DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), - DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64, torch.int64)), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.int64,)), ), decorators=( # Strides are not the same! @@ -17618,6 +17614,9 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k 'TestUnaryUfuncs', "test_reference_numerics_extremal", dtypes=(torch.complex64, torch.complex128), device_type='cpu', active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), ), # tan(j * pi/2 * odd_number) is nan which also make tanhshrink nan. reference_numerics_filter=NumericsFilter( @@ -18002,30 +18001,15 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k sample_inputs_func=sample_inputs_mode,), make_mvlgamma_opinfo(variant_test_name='mvlgamma_p_1', domain=(1, None), - skips=skips_mvlgamma() + ( - # NotImplementedError: The operator 'aten::mvlgamma.out' is not currently - # implemented for the MPS device - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), - ), + skips=skips_mvlgamma(), sample_kwargs=lambda device, dtype, input: ({'p': 1}, {'d': 1})), make_mvlgamma_opinfo(variant_test_name='mvlgamma_p_3', domain=(2, None), - skips=skips_mvlgamma() + ( - # NotImplementedError: The operator 'aten::mvlgamma.out' is not currently - # implemented for the MPS device - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), - ), + skips=skips_mvlgamma(), sample_kwargs=lambda device, dtype, input: ({'p': 3}, {'d': 3})), make_mvlgamma_opinfo(variant_test_name='mvlgamma_p_5', domain=(3, None), - skips=skips_mvlgamma() + ( - # NotImplementedError: The operator 'aten::mvlgamma.out' is not currently - # implemented for the MPS device - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), - ), + skips=skips_mvlgamma(), sample_kwargs=lambda device, dtype, input: ({'p': 5}, {'d': 5})), BinaryUfuncInfo('ne', ref=np.not_equal, @@ -18432,6 +18416,10 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', active_if=IS_WINDOWS), DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=(torch.cfloat, torch.cdouble), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS), DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=2e-3)}), @@ -18464,6 +18452,11 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS), DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=(IS_MACOS or IS_WINDOWS)), @@ -18802,6 +18795,11 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS), DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), # FIXME: @@ -18851,6 +18849,11 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS), DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), DecorateInfo(toleranceOverride({torch.complex64: tol(atol=3e-5, rtol=7e-6)}), @@ -19098,6 +19101,8 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k # tensor(inf+nanj, device='cuda:0') DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cuda', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble]), DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_meta_inplace', dtypes=[torch.bool]), DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_meta_inplace', @@ -19558,10 +19563,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k supports_forward_ad=True, supports_fwgrad_bwgrad=True, error_inputs_func=error_inputs_gather, - skips=( - # RuntimeError: gather(): Yet not supported for complex - DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), - ), ), OpInfo('index_fill', dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.complex32), @@ -19766,9 +19767,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k # Compiler issue on ROCm. Regression started in ROCm 6.4. DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', dtypes=[torch.bool], active_if=TEST_WITH_ROCM), - # RuntimeError: scatter(): Yet not supported for complex - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), - DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), )), UnaryUfuncInfo( 'bfloat16', @@ -20582,9 +20580,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k # Compiler issue on ROCm. Regression started in ROCm 6.4. DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', dtypes=[torch.bool], active_if=TEST_WITH_ROCM), - # RuntimeError: scatter(): Yet not supported for complex - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), - DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), )), OpInfo('stack', dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), @@ -20846,9 +20841,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k sample_inputs_func=sample_repeat_tile, skips=( DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), - # Exception: repeat(): Not supported for complex yet! - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), - DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), )), OpInfo('squeeze', ref=_squeeze_ref, @@ -20962,9 +20954,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k decorators=( # RuntimeError: view size is not compatible with input tensor's size and stride DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides"), - # MPS: gather(): Yet not supported for complex - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), - DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), )), ShapeFuncInfo('tile', ref=np.tile, @@ -20974,11 +20963,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k supports_out=False, supports_forward_ad=True, supports_fwgrad_bwgrad=True, - skips=( - # RuntimeError: repeat(): Not supported for complex yet! - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), - DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), - ), sample_inputs_func=sample_repeat_tile), OpInfo('trapz', # TODO: in the future, 'trapz' should be made a proper alias of 'trapezoid' dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), @@ -21020,16 +21004,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k toleranceOverride({torch.float16: tol(atol=4e-3, rtol=4e-3)}), 'TestInductorOpInfo', 'test_comprehensive', ), - # RuntimeError: cumulative ops are not yet supported for complex - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', - device_type='mps', dtypes=(torch.complex64,) - ), - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', - device_type='mps', dtypes=(torch.complex64,) - ), - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), ), sample_inputs_func=sample_cumulative_trapezoid,), OpInfo('unsqueeze', @@ -21330,8 +21304,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k "test_comprehensive", device_type="cuda" ), - # RuntimeError: Failed to create function state object for: logcumsumexp_outer_float2 - DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), # AssertionError: UserWarning not triggered : Resized a non-empty tensor but did not warn about it. DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), ), @@ -21589,7 +21561,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k sample_inputs_func=sample_inputs_nonzero_static, supports_out=False, supports_autograd=False, - decorators=[onlyCPU], skips=( DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), @@ -22187,7 +22158,7 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k supports_out=False, sample_inputs_func=sample_inputs_grid_sample, reference_inputs_func=reference_inputs_grid_sample, - supports_gradgrad=False, + supports_gradgrad=True, skips=( # Exception: The operator 'aten::grid_sampler_2d_backward' is not currently implemented for the MPS device DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), @@ -22201,7 +22172,7 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k dtypes=floating_types_and(torch.float16, torch.bfloat16), supports_out=False, sample_inputs_func=sample_inputs_grid_sampler_2d, - supports_gradgrad=False, + supports_gradgrad=True, gradcheck_nondet_tol=1e-15, skips=( DecorateInfo(slowTest, 'TestDecomp', 'test_comprehensive', dtypes=(torch.float32, torch.float64), @@ -22218,7 +22189,7 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k dtypes=floating_types_and(torch.float16, torch.bfloat16), supports_out=False, sample_inputs_func=sample_inputs_grid_sampler_3d, - supports_gradgrad=False, + supports_gradgrad=True, gradcheck_nondet_tol=1e-15, skips=( # NOTE: Only run on MPS @@ -22640,8 +22611,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', device_type='xpu', dtypes=[torch.complex128, torch.int8, torch.int16, torch.int32, torch.int64]), - # RuntimeError: cumulative ops are not yet supported for complex - DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), ), ), ReductionOpInfo( @@ -23692,9 +23661,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref'), - # MPS: gather(): Yet not supported for complex - DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), - DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), ), ), PythonRefInfo( @@ -23799,6 +23765,10 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k 'test_reference_numerics_large', device_type='cuda', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), ), ), ElementwiseUnaryPythonRefInfo( @@ -23823,6 +23793,10 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k 'test_reference_numerics_large', device_type='cuda', dtypes=[torch.cfloat], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), ), ), ElementwiseUnaryPythonRefInfo( @@ -23875,6 +23849,10 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k 'test_reference_numerics_large', device_type='cuda', dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), ), ), ElementwiseUnaryPythonRefInfo( @@ -23909,6 +23887,10 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k 'test_reference_numerics_large', device_type='cuda', dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), ), ), ElementwiseUnaryPythonRefInfo( @@ -23953,6 +23935,10 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', dtypes=[torch.chalf], + active_if=IS_WINDOWS), ), ), ElementwiseUnaryPythonRefInfo( @@ -23970,6 +23956,10 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), ), ), ElementwiseUnaryPythonRefInfo( @@ -24279,6 +24269,14 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k 'test_reference_numerics_large', dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), ), ), ElementwiseUnaryPythonRefInfo( @@ -24313,6 +24311,14 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', device_type='cpu', dtypes=[torch.int8]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), ), ), PythonRefInfo( @@ -24386,6 +24392,14 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k 'test_reference_numerics_large', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), ) ), ElementwiseUnaryPythonRefInfo( @@ -24405,6 +24419,14 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k 'test_reference_numerics_large', device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), ), ), ElementwiseUnaryPythonRefInfo( @@ -24852,6 +24874,10 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k dtypes=(torch.complex64, torch.complex128), device_type='cpu', active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), ), ), ElementwiseUnaryPythonRefInfo( @@ -25936,11 +25962,12 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k skips=( # RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(torch.int32, torch.int16) + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.int32, torch.int16, torch.int8, torch.uint8) ), DecorateInfo( unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps', - dtypes=(torch.int32, torch.int16) + dtypes=(torch.int32, torch.int16, torch.int8, torch.uint8) ), ) ), @@ -26003,17 +26030,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k "_refs.repeat", torch_opinfo_name="repeat", validate_view_consistency=False, - skips=( - # Exception: repeat(): Not supported for complex yet! - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_python_ref', - device_type='mps', dtypes=(torch.complex64,) - ), - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', - device_type='mps', dtypes=(torch.complex64,) - ), - ), ), PythonRefInfo( "_refs.reshape", @@ -26366,15 +26382,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k skips=( # doesn't test out behavior properly for this operator DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), - # RuntimeError: cumulative ops are not yet supported for complex - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', - dtypes=(torch.complex64,) - ), - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', - dtypes=(torch.complex64,) - ), ), ), PythonRefInfo( @@ -26384,15 +26391,6 @@ def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **k skips=( # doesn't test out behavior properly for this operator DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), - # RuntimeError: cumulative ops are not yet supported for complex - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', - dtypes=(torch.complex64,) - ), - DecorateInfo( - unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', - dtypes=(torch.complex64,) - ), ), ), PythonRefInfo( diff --git a/torch/testing/_internal/common_modules.py b/torch/testing/_internal/common_modules.py index 57303de39c4e7..fb8202c21450e 100644 --- a/torch/testing/_internal/common_modules.py +++ b/torch/testing/_internal/common_modules.py @@ -2285,6 +2285,47 @@ def module_inputs_torch_nn_MaxPool2d(module_info, device, dtype, requires_grad, desc='return_indices'), ] + +def module_error_inputs_torch_nn_MaxPool2d(module_info, device, dtype, requires_grad, training, **kwargs): + """ + Error inputs for MaxPool2d that test error messages for invalid inputs. + """ + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + # Wrong input dimensions: 2D input instead of 3D/4D + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(2), + forward_input=FunctionInput(make_input((3, 4))), # 2D input + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex=r"non-empty 3D or 4D \(batch mode\) tensor expected for input" + ), + # Wrong input dimensions: 5D input + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(2), + forward_input=FunctionInput(make_input((1, 2, 3, 4, 5))), # 5D input + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex=r"non-empty 3D or 4D \(batch mode\) tensor expected for input" + ), + # Invalid padding: padding > kernel_size / 2 + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(3, padding=5), # kernel=3, pad=5 > 3/2 + forward_input=FunctionInput(make_input((1, 1, 10, 10))), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex=r"pad should be at most half of effective kernel size" + ), + ] + + def module_inputs_torch_nn_MaxPool3d(module_info, device, dtype, requires_grad, training, **kwargs): make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) @@ -2844,6 +2885,57 @@ def module_inputs_torch_nn_Embedding(module_info, device, dtype, requires_grad, ] +def module_error_inputs_torch_nn_Embedding(module_info, device, dtype, requires_grad, training, **kwargs): + """ + Error inputs for Embedding that test error messages for invalid inputs. + """ + samples = [] + + # Out of range indices: index exceeds num_embeddings + # Only test on CPU - CUDA triggers kernel assertion instead of Python exception + if torch.device(device).type == 'cpu': + samples.append( + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(num_embeddings=10, embedding_dim=3), + forward_input=FunctionInput(torch.tensor([0, 5, 15], device=device, dtype=torch.long)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=IndexError, + error_regex=r"index out of range in self" + ) + ) + + # Float indices: wrong dtype for indices (works on all devices) + samples.append( + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(num_embeddings=10, embedding_dim=3), + forward_input=FunctionInput(torch.tensor([1.5, 2.5], device=device, dtype=torch.float32)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex=r"Expected tensor for argument.*indices.*to have.*scalar type.*Long.*Int" + ) + ) + + # Negative num_embeddings (construction error, device-independent) + samples.append( + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(num_embeddings=-1, embedding_dim=3), + forward_input=FunctionInput(), + ), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=RuntimeError, + error_regex=r"Trying to create tensor with negative dimension" + ) + ) + + return samples + + + def module_inputs_torch_nn_MultiheadAttention(module_info, device, dtype, requires_grad, training, **kwargs): # Currently all samples below are for validating the no-batch-dim support. make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) @@ -4103,6 +4195,7 @@ def module_error_inputs_torch_nn_Pad3d(module_info, device, dtype, requires_grad ), ModuleInfo(torch.nn.MaxPool2d, module_inputs_func=module_inputs_torch_nn_MaxPool2d, + module_error_inputs_func=module_error_inputs_torch_nn_MaxPool2d, ), ModuleInfo(torch.nn.MaxPool3d, module_inputs_func=module_inputs_torch_nn_MaxPool3d, @@ -4389,6 +4482,7 @@ def module_error_inputs_torch_nn_Pad3d(module_info, device, dtype, requires_grad ), ModuleInfo(torch.nn.Embedding, module_inputs_func=module_inputs_torch_nn_Embedding, + module_error_inputs_func=module_error_inputs_torch_nn_Embedding, gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, decorators=[ DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-4)}), diff --git a/torch/testing/_internal/common_mps.py b/torch/testing/_internal/common_mps.py index 0e7d23468737e..f8fc0f4437d8b 100644 --- a/torch/testing/_internal/common_mps.py +++ b/torch/testing/_internal/common_mps.py @@ -55,6 +55,9 @@ def mps_ops_modifier( "cos", "cosh", "cross", + "cumsum", + "cumprod", + "cumulative_trapezoid", "diag", "diag_embed", "diagflat", @@ -73,6 +76,7 @@ def mps_ops_modifier( "expand", "expand_as", "expand_copy", + "gather", "flatten", "fill", "full", @@ -95,6 +99,7 @@ def mps_ops_modifier( "linalg.diagonal", "linalg.householder_product", "linalg.svd", + "linalg.vander", "linalg.vecdot", "linalg.vector_norm", "log10", @@ -103,6 +108,7 @@ def mps_ops_modifier( "log", "logaddexp", "logaddexp2", + "logcumsumexp", "mH", "mT", "masked_fill", @@ -125,6 +131,7 @@ def mps_ops_modifier( "nn.functional.conv_transpose3d", "nn.functional.feature_alpha_dropoutwithout_train", "nn.functional.l1_loss", + "nn.functional.linear", "nn.functional.normalize", "nn.functional.padcircular", "nn.functional.pairwise_distance", @@ -135,6 +142,7 @@ def mps_ops_modifier( "nn.functional.triplet_margin_with_distance_loss", "nn.functional.unfold", "nonzero", + "nonzero_static", "norm", "normfro", "norminf", @@ -147,6 +155,7 @@ def mps_ops_modifier( "randn", "ravel", "real", + "repeat", "repeat_interleave", "reshape_as", "reshape", @@ -155,6 +164,8 @@ def mps_ops_modifier( "rsqrt", "rsub", "scalar_tensor", + "scatter", + "scatter_add", "select", "sgn", "sigmoid", @@ -179,9 +190,11 @@ def mps_ops_modifier( "svd", "t", "t_copy", + "take_along_dim", "tanh", "tan", "tensor_split", + "tile", "transpose", "transpose_copy", "tril", @@ -275,6 +288,8 @@ def mps_ops_modifier( "logical_xor", "logsumexp", "long", + "masked.cumsum", + "masked.cumprod", "masked.mean", "masked.normalize", "masked.prod", @@ -618,7 +633,6 @@ def mps_ops_modifier( "float_power": None, "linalg.matrix_rankhermitian": None, "linalg.pinvhermitian": None, - "nonzero_static": None, # MPS: input sizes must be divisible by output sizes "nn.functional.adaptive_avg_pool1d": None, "nn.functional.adaptive_avg_pool2d": None, @@ -916,7 +930,6 @@ def mps_ops_grad_modifier(ops: Sequence[OpInfo]) -> Sequence[OpInfo]: "grid_sampler_3d": None, "igamma": None, # currently not supported for any device "igammac": None, # currently not supported for any device - "aminmax": [torch.float32, torch.float16], "special.i1": [torch.float16], # "i1_backward" not implemented for 'Half' "special.i1e": [torch.float16], # "i1e_backward" not implemented for 'Half' # Correctness issues diff --git a/torch/testing/_internal/common_ops_unbacked.py b/torch/testing/_internal/common_ops_unbacked.py index 56a47a263cfc4..766126de04375 100644 --- a/torch/testing/_internal/common_ops_unbacked.py +++ b/torch/testing/_internal/common_ops_unbacked.py @@ -124,7 +124,6 @@ def skip(op_name, variant_name="", *, device_type=None, dtypes=None): xfail("masked.var"), xfail("max_pool2d_with_indices_backward"), xfail("multinomial"), - xfail("nanquantile"), xfail("nn.functional.adaptive_avg_pool1d"), xfail("nn.functional.adaptive_avg_pool2d"), xfail("nn.functional.adaptive_avg_pool3d"), @@ -153,7 +152,6 @@ def skip(op_name, variant_name="", *, device_type=None, dtypes=None): xfail("nn.functional.fractional_max_pool2d"), xfail("nn.functional.fractional_max_pool3d"), xfail("nn.functional.gaussian_nll_loss"), - xfail("nn.functional.glu"), xfail("nn.functional.grid_sample"), xfail("nn.functional.group_norm"), xfail("nn.functional.huber_loss"), @@ -196,7 +194,6 @@ def skip(op_name, variant_name="", *, device_type=None, dtypes=None): xfail("ormqr"), xfail("pca_lowrank"), xfail("pinverse"), - xfail("quantile"), xfail("qr"), xfail("rand_like"), xfail("randint_like"), diff --git a/torch/testing/_internal/common_optimizers.py b/torch/testing/_internal/common_optimizers.py index 6d0ec90052684..fb67c4d845021 100644 --- a/torch/testing/_internal/common_optimizers.py +++ b/torch/testing/_internal/common_optimizers.py @@ -506,6 +506,28 @@ def optim_error_inputs_func_adagrad(device, dtype): error_regex="Invalid lr_decay value: -0.5", ), ] + if _get_device_type(device) == "cuda": + sample_tensor = torch.empty((), device=device, dtype=dtype) + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=[sample_tensor], + kwargs={"foreach": True, "fused": True}, + desc="`fused` and `foreach` cannot be `True` together", + ), + error_type=RuntimeError, + error_regex="`fused` and `foreach` cannot be `True` together", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[sample_tensor], + kwargs={"fused": True, "differentiable": True}, + desc="`fused` does not support `differentiable`", + ), + error_type=RuntimeError, + error_regex="`fused` does not support `differentiable`", + ), + ] return error_inputs @@ -1599,7 +1621,7 @@ def _get_optim_inputs_including_global_cliquey_kwargs( "maximize", "capturable", ), - supports_fused_on=("cpu",), + supports_fused_on=("cpu", "cuda"), supports_sparse=True, metadata_for_sparse=( {"lr": 0.1, "weight_decay": 0, "lr_decay": 0}, @@ -1624,6 +1646,16 @@ def _get_optim_inputs_including_global_cliquey_kwargs( "TestOptimRenewed", "test_fused_matches_forloop", ), + DecorateInfo( + toleranceOverride( + { # https://github.com/pytorch/pytorch/issues/116202 + torch.float32: tol(atol=5e-04, rtol=0.015), + } + ), + "TestOptimRenewed", + "test_mixed_device_dtype", + active_if=TEST_WITH_TORCHDYNAMO, + ), ), skips=( DecorateInfo( diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py index 8b8507433c7d8..193d625c5e9ce 100644 --- a/torch/testing/_internal/common_utils.py +++ b/torch/testing/_internal/common_utils.py @@ -82,17 +82,19 @@ ) from torch.testing import make_tensor from torch.testing._comparison import ( + _unwrap_dtensor_for_comparison, BooleanPair, NonePair, + not_close_error_metas, NumberPair, Pair, TensorLikePair, ) -from torch.testing._comparison import not_close_error_metas from torch.testing._internal.common_dtype import get_all_dtypes from torch.utils._import_utils import _check_module_exists import torch.utils._pytree as pytree from torch.utils import cpp_extension +from torch._utils import _is_privateuse1_backend_available try: import pytest # type: ignore[import-not-found] has_pytest = True @@ -1452,12 +1454,6 @@ def TemporaryDirectoryName(suffix=None): yield d -def is_privateuse1_backend_available(): - privateuse1_backend_name = torch._C._get_privateuse1_backend_name() - privateuse1_backend_module = getattr(torch, privateuse1_backend_name, None) - return (is_available := getattr(privateuse1_backend_module, "is_available", None)) and is_available() - - def make_lazy_class(cls): def lazy_init(self, cb): @@ -1511,7 +1507,7 @@ class LazyVal: TEST_ACCELERATOR = LazyVal(lambda: torch.accelerator.is_available()) # type: ignore[call-arg] TEST_MULTIACCELERATOR = LazyVal(lambda: torch.accelerator.device_count() > 1) # type: ignore[call-arg] custom_device_mod = getattr(torch, torch._C._get_privateuse1_backend_name(), None) -TEST_PRIVATEUSE1 = is_privateuse1_backend_available() +TEST_PRIVATEUSE1 = _is_privateuse1_backend_available() TEST_PRIVATEUSE1_DEVICE_TYPE = torch._C._get_privateuse1_backend_name() TEST_NUMBA = _check_module_exists('numba') TEST_TRANSFORMERS = _check_module_exists('transformers') @@ -1523,6 +1519,44 @@ class LazyVal: TEST_Z3 = _check_module_exists('z3') +# DSL availability (lazy evaluation to avoid import overhead) +class LazyDSLCheck: + """Lazy DSL availability checker to avoid import-time overhead""" + def __init__(self): + self._registry = None + self._import_attempted = False + + def _get_registry(self): + if not self._import_attempted: + self._import_attempted = True + try: + from torch._native.dsl_registry import dsl_registry + self._registry = dsl_registry + except ImportError: + self._registry = None + return self._registry + + def is_available(self, dsl_name: str) -> bool: + """Check if specific DSL is available""" + registry = self._get_registry() + return registry.is_dsl_available(dsl_name) if registry is not None else False + + def list_available(self) -> list[str]: + """Get list of available DSLs""" + registry = self._get_registry() + return list(registry.list_available_dsls()) if registry is not None else [] + + def list_all(self) -> list[str]: + """Get list of all registered DSLs""" + registry = self._get_registry() + return list(registry.list_all_dsls()) if registry is not None else [] + +_dsl_checker = LazyDSLCheck() + +# Lazy constants to avoid import-time overhead +TEST_TRITON_DSL = LazyVal(lambda: _dsl_checker.is_available('triton')) +TEST_CUTEDSL = LazyVal(lambda: _dsl_checker.is_available('cutedsl')) + def split_if_not_empty(x: str): return x.split(",") if len(x) != 0 else [] @@ -1530,6 +1564,34 @@ def split_if_not_empty(x: str): skipIfNoDill = unittest.skipIf(not TEST_DILL, "no dill") +# DSL skip decorators (following existing pattern) +skipIfNoTritonDSL = unittest.skipIf(not TEST_TRITON_DSL, "Triton DSL not available") +skipIfNoCuteDSL = unittest.skipIf(not TEST_CUTEDSL, "CuTeDSL not available") + +def skipIfDSLUnavailable(dsl_name: str, reason: str | None = None): + """Skip test if specific DSL is not available""" + available = _dsl_checker.is_available(dsl_name) + msg = reason or f"{dsl_name} DSL not available" + return unittest.skipIf(not available, msg) + +def skipUnlessDSLAvailable(dsl_name: str, reason: str | None = None): + """Skip test unless specific DSL is available""" + available = _dsl_checker.is_available(dsl_name) + msg = reason or f"{dsl_name} DSL required" + return unittest.skipUnless(available, msg) + +def get_available_dsls() -> list[str]: + """Get list of available DSL names for test parameterization""" + return _dsl_checker.list_available() + +def is_dsl_available(dsl_name: str) -> bool: + """Check if specific DSL is available for conditional testing""" + return _dsl_checker.is_available(dsl_name) + +def get_all_dsls() -> list[str]: + """Get all registered DSL names (available or not) for comprehensive testing""" + return _dsl_checker.list_all() + NO_MULTIPROCESSING_SPAWN: bool = False TEST_WITH_ASAN: bool = TestEnvironment.def_flag( @@ -1551,6 +1613,7 @@ def split_if_not_empty(x: str): TEST_WITH_ROCM: bool = TestEnvironment.def_flag( "TEST_WITH_ROCM", env_var="PYTORCH_TEST_WITH_ROCM", + implied_by_fn=lambda: torch.version.hip is not None, ) TEST_WITH_MTIA: bool = TestEnvironment.def_flag( "TEST_WITH_MTIA", @@ -2002,7 +2065,7 @@ def has_corresponding_torch_dtype(np_dtype): def skipIfNNModuleInlined( msg="test doesn't currently work with nn module inlining", - condition=torch._dynamo.config.inline_inbuilt_nn_modules, + condition=True, ): def decorator(fn): if not isinstance(fn, type): @@ -2108,12 +2171,31 @@ def wrapper(*args, **kwargs): return dec_fn def skipIfMPS(fn): + sig = inspect.signature(fn) + has_device_arg = "device" in sig.parameters + + if not has_device_arg: + warnings.warn( + f"skipIfMPS applied to {fn.__qualname__} which has no 'device' parameter. " + "Consider using device-generic tests with instantiate_device_type_tests instead.", + stacklevel=2, + ) + @wraps(fn) def wrapper(*args, **kwargs): - if TEST_MPS: + if has_device_arg: + # For device-generic tests, only skip when actually running on MPS + slf = args[0] if args else None + if slf is not None: + device_type = getattr(slf, "device_type", None) or getattr( + slf, "device", None + ) + if isinstance(device_type, str) and device_type == "mps": + raise unittest.SkipTest("test doesn't currently work with MPS") + elif TEST_MPS: raise unittest.SkipTest("test doesn't currently work with MPS") - else: - fn(*args, **kwargs) + return fn(*args, **kwargs) + return wrapper @@ -4315,6 +4397,8 @@ def to_list(input): if isinstance(y, torch.Tensor) and y.is_nested and y.layout == torch.strided: y = y.unbind() + x, y = _unwrap_dtensor_for_comparison(x, y) + error_metas = not_close_error_metas( x, y, diff --git a/torch/testing/_internal/common_xpu.py b/torch/testing/_internal/common_xpu.py new file mode 100644 index 0000000000000..ee4f8c921eac5 --- /dev/null +++ b/torch/testing/_internal/common_xpu.py @@ -0,0 +1,88 @@ +import enum +import functools + +import torch +import torch.xpu +from torch.testing._internal.common_utils import IS_WINDOWS, LazyVal, TEST_XPU + + +XPU_ALREADY_INITIALIZED_ON_IMPORT = torch.xpu.is_initialized() + + +class XPUCodename(enum.Enum): + PVC = "PVC" # Intel® Data Center GPU Max Series + BMG = "BMG" # Intel® Arc™ Pro Battlemage Graphics + + +class XPUArch(enum.IntEnum): + Unknown = 0 + Xe = 1 # Xe HPC + Xe2 = 2 + + +# device_id -> GPU codename +# From https://github.com/intel/intel-graphics-compiler/blob/master/inc/common/igfxfmid.h +_DEVICE_ID_TO_CODENAME = { + 0x0BD0: XPUCodename.PVC, + 0x0BD4: XPUCodename.PVC, + 0x0BD5: XPUCodename.PVC, + 0x0BD6: XPUCodename.PVC, + 0x0BD7: XPUCodename.PVC, + 0x0BD8: XPUCodename.PVC, + 0x0BD9: XPUCodename.PVC, + 0x0BDA: XPUCodename.PVC, + 0x0BDB: XPUCodename.PVC, + 0x0B69: XPUCodename.PVC, + 0x0B6E: XPUCodename.PVC, + 0xE202: XPUCodename.BMG, + 0xE20B: XPUCodename.BMG, + 0xE20C: XPUCodename.BMG, + 0xE20D: XPUCodename.BMG, + 0xE210: XPUCodename.BMG, + 0xE212: XPUCodename.BMG, + 0xE215: XPUCodename.BMG, + 0xE216: XPUCodename.BMG, + 0xE220: XPUCodename.BMG, + 0xE221: XPUCodename.BMG, + 0xE222: XPUCodename.BMG, + 0xE223: XPUCodename.BMG, +} + +# GPU codename -> architecture +_CODENAME_TO_ARCH = { + XPUCodename.PVC: XPUArch.Xe, + XPUCodename.BMG: XPUArch.Xe2, +} + + +@functools.lru_cache(1) +def get_xpu_codename() -> XPUCodename | None: + device_id = torch.xpu.get_device_capability()["device_id"] + return _DEVICE_ID_TO_CODENAME.get(device_id) + + +@functools.lru_cache(1) +def get_xpu_arch() -> XPUArch | None: + codename = get_xpu_codename() + return _CODENAME_TO_ARCH.get(codename, XPUArch.Unknown) + + +Xe2_Or_Later = LazyVal( + lambda: torch.xpu.is_available() and get_xpu_arch() >= XPUArch.Xe2 +) + + +def evaluate_platform_supports_flash_attention(): + if TEST_XPU: + return not IS_WINDOWS and Xe2_Or_Later + return False + + +PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU: bool = LazyVal( + lambda: evaluate_platform_supports_flash_attention() +) + +# Importing this module should NOT eagerly initialize XPU +if not XPU_ALREADY_INITIALIZED_ON_IMPORT: + if torch.xpu.is_initialized(): + raise AssertionError("XPU should not be initialized on import") diff --git a/torch/testing/_internal/distributed/_tensor/common_dtensor.py b/torch/testing/_internal/distributed/_tensor/common_dtensor.py index dde1b4d7bc491..68c6baa863e98 100644 --- a/torch/testing/_internal/distributed/_tensor/common_dtensor.py +++ b/torch/testing/_internal/distributed/_tensor/common_dtensor.py @@ -1461,7 +1461,49 @@ def _to_local_shard(a): dt = DTensor.from_local(local, device_mesh, (plc,)) full = dt.redistribute(device_mesh, (Replicate(),)).to_local() if ref.shape != full.shape or not torch.allclose( - ref, full, atol=1e-5, rtol=1e-5 + ref, full, atol=1e-5, rtol=1e-5, equal_nan=True ): return False return True + + +@contextlib.contextmanager +def op_strategy_context(op_overload, strategy_func, schema_info=None): + """ + Context manager for setting and clearing op strategies. + Args: + op_overload: The operator overload to set or clear the strategy for. + strategy_func: The strategy function to set for the operator overload. + schema_info: Optional schema information for the operator overload. + Yields: + None + """ + from torch.distributed.tensor._ops.utils import register_op_strategy + from torch.distributed.tensor.debug import _clear_sharding_prop_cache + + propagator = DTensor._op_dispatcher.sharding_propagator + _origin_op_strategy_funcs = None + _origin_op_strategy_schema = None + try: + # register the op strategy + if op_overload in propagator.op_strategy_funcs: + _origin_op_strategy_funcs = propagator.op_strategy_funcs[op_overload] + del propagator.op_strategy_funcs[op_overload] + if op_overload in propagator.op_to_schema_info: + _origin_op_strategy_schema = propagator.op_to_schema_info[op_overload] + del propagator.op_to_schema_info[op_overload] + register_op_strategy(op_overload, schema_info=schema_info)(strategy_func) + yield + finally: + # clear this op strategy cache + if _origin_op_strategy_funcs is None: + if op_overload in propagator.op_strategy_funcs: + del propagator.op_strategy_funcs[op_overload] + else: + propagator.op_strategy_funcs[op_overload] = _origin_op_strategy_funcs + if _origin_op_strategy_schema is None: + if op_overload in propagator.op_to_schema_info: + del propagator.op_to_schema_info[op_overload] + else: + propagator.op_to_schema_info[op_overload] = _origin_op_strategy_schema + _clear_sharding_prop_cache() diff --git a/torch/testing/_internal/distributed/distributed_test.py b/torch/testing/_internal/distributed/distributed_test.py index f182a724773ab..3f96956170976 100644 --- a/torch/testing/_internal/distributed/distributed_test.py +++ b/torch/testing/_internal/distributed/distributed_test.py @@ -198,9 +198,6 @@ class TestNamedTupleInput_1(NamedTuple): BACKEND = os.environ["BACKEND"] INIT_METHOD = os.getenv("INIT_METHOD", "env://") -DEFAULT_TIMEOUT = 300 -CUSTOMIZED_TIMEOUT = {"test_DistributedDataParallel": 500} - def get_profiling_event(event_name, profiler, dedup_gpu_user_annotation=False): event_list = ( @@ -402,14 +399,6 @@ def forward(self, x): return F.relu(self.lin1(x)) -def get_timeout(test_id): - test_name = test_id.split(".")[-1] - if test_name in CUSTOMIZED_TIMEOUT: - return CUSTOMIZED_TIMEOUT[test_name] - else: - return DEFAULT_TIMEOUT - - default_pg_timeout = 60 CUSTOM_PG_TIMEOUT = { @@ -905,6 +894,24 @@ def test_barrier_timeout_full_group(self): if group_id is not None: self._test_barrier_timeout(group_id, timeout) + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo", "Only gloo backend supports timeouts" + ) + def test_barrier_timeout_arg(self): + """Test that the timeout argument to barrier() overrides PG default. + + Create a PG with a large default timeout, then have only rank 0 + call barrier with a tiny timeout. The barrier should time out using + the per-call timeout (1ms) rather than the PG default (300s). + """ + pg = dist.new_group(timeout=timedelta(seconds=300)) + + if dist.get_rank() == 0: + with self.assertRaisesRegex(RuntimeError, "Timed out waiting 1ms"): + dist.barrier(group=pg, timeout=timedelta(seconds=0.001)) + + dist.destroy_process_group(pg) + @skip_but_pass_in_sandcastle_if( BACKEND not in DistTestCases.backend_feature["subgroup"], f"The {BACKEND} backend does not support creating subgroups on CUDA devices", diff --git a/torch/testing/_internal/hop_db.py b/torch/testing/_internal/hop_db.py index d4b03757a0877..76e22c62e27ff 100644 --- a/torch/testing/_internal/hop_db.py +++ b/torch/testing/_internal/hop_db.py @@ -5,7 +5,14 @@ import torch from functorch.experimental.control_flow import map -from torch.nn.attention.flex_attention import _create_empty_block_mask, flex_attention +from torch._higher_order_ops.flex_attention import ( + flex_attention as flex_attention_hop, +) +from torch.nn.attention.flex_attention import ( + _create_empty_block_mask, + create_block_mask, + flex_attention, +) from torch.testing import make_tensor from torch._higher_order_ops.inline_asm_elementwise import inline_asm_elementwise from torch.testing._internal.common_device_type import onlyCUDA @@ -196,6 +203,99 @@ def score_mod(score, b, h, m, n): yield SampleInput(q, k, v, score_mod, block_mask) +def sample_inputs_flex_attention_backward( + opinfo, device, dtype, requires_grad, **kwargs +): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=False + ) + + def score_mod(score, b, h, m, n): + return score + + def mask_mod(b, h, m, n): + return m >= n + + q, k, v = (make_arg(2, 2, 128, 16, low=0.1, high=2) for _ in range(3)) + block_mask = create_block_mask(mask_mod, B=2, H=2, Q_LEN=128, KV_LEN=128, device=device) + scale = 1.0 / q.size(-1) ** 0.5 + out, logsumexp, _ = flex_attention_hop( + q, k, v, score_mod, block_mask.as_tuple(), scale, {}, + ) + yield SampleInput( + q, + args=( + k, v, out.detach(), logsumexp.detach(), torch.rand_like(out), None, + score_mod, None, block_mask.as_tuple(), + scale, {}, (), (), + ), + ) + + +def sample_inputs_flex_attention_backward_explicit_buffers( + opinfo, device, dtype, requires_grad, **kwargs +): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=False + ) + mask_offset = torch.full((), 128, device=device, dtype=torch.int32) + + def score_mod(score, b, h, m, n): + return score + + def mask_mod(b, h, m, n): + return m + mask_offset >= n + + q, k, v = (make_arg(2, 2, 128, 16, low=0.1, high=2) for _ in range(3)) + block_mask = create_block_mask(mask_mod, B=2, H=2, Q_LEN=128, KV_LEN=128, device=device) + scale = 1.0 / q.size(-1) ** 0.5 + out, logsumexp, _ = flex_attention_hop( + q, k, v, score_mod, block_mask.as_tuple(), scale, {}, + ) + yield SampleInput( + q, + args=( + k, v, out.detach(), logsumexp.detach(), torch.rand_like(out), None, + score_mod, None, block_mask.as_tuple(), + scale, {}, (), (), + ), + ) + + +def simple_flex_attention_backward( + query, + key, + value, + out, + logsumexp, + grad_out, + grad_logsumexp, + fw_graph, + joint_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, +): + return torch.ops.higher_order.flex_attention_backward( + query, + key, + value, + out, + logsumexp, + grad_out, + grad_logsumexp, + fw_graph, + joint_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + + def sample_inputs_while_loop(opinfo, device, dtype, requires_grad, **kwargs): make_arg = functools.partial( make_tensor, device=device, dtype=dtype, requires_grad=False @@ -490,14 +590,39 @@ def simple_inline_asm(x): OpInfo( name="flex_attention_backward", variant_test_name="simple", - op=flex_attention, - sample_inputs_func=sample_inputs_flex_attention, + op=simple_flex_attention_backward, + sample_inputs_func=sample_inputs_flex_attention_backward, dtypes=custom_types(torch.float16, torch.float32), supports_out=False, check_batched_grad=False, check_batched_gradgrad=False, check_batched_forward_grad=False, check_inplace_batched_forward_grad=False, + supports_autograd=False, + supports_gradgrad=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"), + DecorateInfo( + unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export" + ), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"), + ), + decorators=[onlyCUDA], + ), + OpInfo( + name="flex_attention_backward", + variant_test_name="explicit_buffers", + op=simple_flex_attention_backward, + sample_inputs_func=sample_inputs_flex_attention_backward_explicit_buffers, + dtypes=custom_types(torch.float16, torch.float32), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=False, + supports_gradgrad=False, skips=( DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"), DecorateInfo( @@ -546,14 +671,6 @@ def simple_inline_asm(x): check_batched_forward_grad=False, check_inplace_batched_forward_grad=False, supports_autograd=False, - skips=( - DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"), - DecorateInfo( - unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export" - ), - DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"), - DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"), - ), decorators=[onlyCUDA], ), ] diff --git a/torch/testing/_internal/opinfo/definitions/_masked.py b/torch/testing/_internal/opinfo/definitions/_masked.py index a1eff576e3765..670b34c64ecdf 100644 --- a/torch/testing/_internal/opinfo/definitions/_masked.py +++ b/torch/testing/_internal/opinfo/definitions/_masked.py @@ -641,16 +641,10 @@ def sample_inputs_masked_normalize(op_info, device, dtype, requires_grad, **kwar DecorateInfo( unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit" ), - # Exception: cumulative ops are not yet supported for complex + # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}. DecorateInfo( unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" ), - DecorateInfo( - unittest.expectedFailure, - "TestCommon", - device_type="mps", - dtypes=(torch.complex64,), - ), ), # Can reuse the same inputs; dim is required in both sample_inputs_func=sample_inputs_masked_cumops, @@ -688,16 +682,10 @@ def sample_inputs_masked_normalize(op_info, device, dtype, requires_grad, **kwar "test_comprehensive", device_type="cuda", ), - # Exception: cumulative ops are not yet supported for complex + # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}. DecorateInfo( unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" ), - DecorateInfo( - unittest.expectedFailure, - "TestCommon", - device_type="mps", - dtypes=(torch.complex64,), - ), ), # Can reuse the same inputs; dim is required in both sample_inputs_func=sample_inputs_masked_cumops, diff --git a/torch/testing/_internal/opinfo/definitions/linalg.py b/torch/testing/_internal/opinfo/definitions/linalg.py index 03fb91cd103aa..b96a66af8db95 100644 --- a/torch/testing/_internal/opinfo/definitions/linalg.py +++ b/torch/testing/_internal/opinfo/definitions/linalg.py @@ -1969,16 +1969,6 @@ def make_input(): "TestCommon", "test_numpy_ref_mps", ), - # Exception: cumulative ops are not yet supported for complex - DecorateInfo( - unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" - ), - DecorateInfo( - unittest.expectedFailure, - "TestCommon", - device_type="mps", - dtypes=(torch.complex64,), - ), ), ), ReductionOpInfo( diff --git a/torch/testing/_internal/opinfo/definitions/sparse.py b/torch/testing/_internal/opinfo/definitions/sparse.py index 7f8781ce87b24..c04009f1f2ac9 100644 --- a/torch/testing/_internal/opinfo/definitions/sparse.py +++ b/torch/testing/_internal/opinfo/definitions/sparse.py @@ -4,6 +4,7 @@ import torch from torch.testing import make_tensor # noqa: F401 +from torch.testing._internal.common_dtype import highest_precision_float from torch.testing._internal.opinfo.core import ( # noqa: F401 BinaryUfuncInfo, ErrorInput, @@ -769,8 +770,9 @@ def _sample_inputs_sparse_like_fns( tensor, args=(), kwargs=dict(device=device, dtype=dtype, layout=layout) ) - if dtype is not torch.float64: - yield SampleInput(tensor, args=(), kwargs=dict(dtype=torch.float64)) + hpf = highest_precision_float(device) + if dtype is not hpf: + yield SampleInput(tensor, args=(), kwargs=dict(dtype=hpf)) if torch.cuda.is_available(): other_device = "cuda" if tensor.device.type == "cpu" else "cpu" diff --git a/torch/utils/_config_module.py b/torch/utils/_config_module.py index 2bb21c31c4c86..7f0a3ec316213 100644 --- a/torch/utils/_config_module.py +++ b/torch/utils/_config_module.py @@ -9,11 +9,11 @@ import tokenize import unittest from collections.abc import Callable +from contextvars import ContextVar from dataclasses import dataclass from types import FunctionType, ModuleType from typing import Any, Generic, NoReturn, Optional, TYPE_CHECKING, TypeVar from typing_extensions import deprecated -from unittest import mock from torch._utils_internal import justknobs_check @@ -46,7 +46,7 @@ class _Config(Generic[T]): If multiple env variables are given, the precedence order is from left to right. user_override: If a user sets a value (i.e. foo.bar=True), that - has precedence over everything after this. + has precedence over everything after this. User overrides are thread-local. env_name_default: If set, this environment variable will override everything after this. If multiple env variables are given, the precedence order is from @@ -203,7 +203,7 @@ def visit( annotated_type = type_hints.get(key, None) if isinstance(value, CONFIG_TYPES): config[name] = _ConfigEntry( - _Config(default=value, value_type=annotated_type) + _Config(default=value, value_type=annotated_type), name ) if dest is module: delattr(module, key) @@ -211,7 +211,7 @@ def visit( if annotated_type is not None and value.value_type is None: value.value_type = annotated_type - config[name] = _ConfigEntry(value) + config[name] = _ConfigEntry(value, name) if dest is module: delattr(module, key) @@ -291,7 +291,7 @@ class _ConfigEntry: value_type: type # The value specified by the user when they overrode the configuration # _UNSET_SENTINEL indicates the value is not set. - user_override: Any = _UNSET_SENTINEL + user_override: ContextVar[object] # The justknob to check for this config justknob: str | None = None # environment variables are read at install time @@ -315,7 +315,7 @@ class _ConfigEntry: deprecation_message: str | None = None _deprecation_warned: bool = False - def __init__(self, config: _Config) -> None: + def __init__(self, config: _Config, name: str) -> None: self.default = config.default self.value_type = ( config.value_type if config.value_type is not None else type(self.default) @@ -327,6 +327,7 @@ def __init__(self, config: _Config) -> None: self.deprecation_message = config.deprecation_message self._deprecation_warned = False + self.user_override = ContextVar(name, default=_UNSET_SENTINEL) if config.env_name_default is not None: for val in config.env_name_default: if (env_value := _read_env_variable(val)) is not None: @@ -400,7 +401,7 @@ def __setattr__(self, name: str, value: object) -> None: if config.alias is not None: self._set_alias_val(config, value) else: - config.user_override = value + config.user_override.set(value) self._is_dirty = True config.hide = False @@ -421,8 +422,9 @@ def __getattr__(self, name: str) -> Any: if config.env_value_force is not _UNSET_SENTINEL: return config.env_value_force - if config.user_override is not _UNSET_SENTINEL: - return config.user_override + user_override = config.user_override.get() + if user_override is not _UNSET_SENTINEL: + return user_override if config.env_value_default is not _UNSET_SENTINEL: return config.env_value_default @@ -435,8 +437,8 @@ def __getattr__(self, name: str) -> Any: # copy them to user_overrides in case the user overrides # them if isinstance(config.default, (list, set, dict)): - config.user_override = copy.deepcopy(config.default) - return config.user_override + config.user_override.set(copy.deepcopy(config.default)) + return config.user_override.get() return config.default except KeyError as e: @@ -447,7 +449,7 @@ def __delattr__(self, name: str) -> None: self._is_dirty = True # must support delete because unittest.mock.patch deletes # then recreate things - self._config[name].user_override = _UNSET_SENTINEL + self._config[name].user_override.set(_UNSET_SENTINEL) self._config[name].hide = True def _get_alias_module_and_name( @@ -498,10 +500,10 @@ def _is_default(self, name: str) -> bool: or config_val.env_value_force == config_val.default ) - unset = config_val.user_override is _UNSET_SENTINEL + unset = config_val.user_override.get() is _UNSET_SENTINEL # Handle reference types specially to avoid spammy warnings if isinstance(config_val.default, (list, set, dict)): - unset = unset or config_val.user_override == config_val.default + unset = unset or config_val.user_override.get() == config_val.default return unset and not_set_env_default and not_set_env_force def _get_dict( @@ -527,7 +529,9 @@ def _get_dict( it skips it. """ config: dict[str, Any] = {} - for key in self._config: + for key, entry in self._config.items(): + if entry.alias is not None: + continue if ignored_keys and key in ignored_keys: continue if ignored_prefixes: @@ -535,14 +539,23 @@ def _get_dict( continue if skip_default and self._is_default(key): continue - if self._config[key].alias is not None: - continue - curr_entry = self._config[key] - has_been_warned = curr_entry._deprecation_warned - curr_entry._deprecation_warned = True - config[key] = copy.deepcopy(getattr(self, key)) - curr_entry._deprecation_warned = has_been_warned + # Read value directly, bypassing __getattr__ overhead + # (deprecation warnings, alias resolution). + user_override = entry.user_override.get() + if entry.env_value_force is not _UNSET_SENTINEL: + val = entry.env_value_force + elif user_override is not _UNSET_SENTINEL: + val = user_override + elif entry.env_value_default is not _UNSET_SENTINEL: + val = entry.env_value_default + elif entry.justknob is not None: + val = justknobs_check(name=entry.justknob, default=entry.default) + else: + val = entry.default + if isinstance(val, (list, set, dict)): + val = copy.deepcopy(val) + config[key] = val return config @@ -565,7 +578,19 @@ def save_config_portable( if ignore_private_configs: prefixes.append("_") prefixes.extend(getattr(self, "_cache_config_ignore_prefix", [])) - return self._get_dict(ignored_prefixes=prefixes) + config = self._get_dict(ignored_prefixes=prefixes) + for key in getattr(self, "_cache_config_factory_keys", []): + if key in config and config[key] is not None: + instance = config[key]() + if hasattr(instance, "uuid"): + config[key] = instance.uuid() + else: + raise RuntimeError( + f"Config '{key}' is set to {config[key]} which does not " + f"implement uuid(). Implement uuid() for cache key " + f"participation." + ) + return config def codegen_config(self) -> str: """Convert config to Python statements that replicate current config. @@ -688,7 +713,7 @@ def patch( **kwargs: dict[str, Any], ) -> "ContextDecorator": """ - Decorator and/or context manager to make temporary changes to a config. + Decorator and/or context manager to make temporary changes to a config. Note that patched settings are thread-local. As a decorator: @@ -732,28 +757,40 @@ def foo(...): ) if not isinstance(changes, dict): raise AssertionError(f"expected `dict` got {type(changes)}") - prior: dict[str, Any] = {} config = self class ConfigPatch(ContextDecorator): def __init__(self) -> None: self.changes = changes + self._prior: ContextVar[tuple[dict[str, Any], ...]] = ContextVar( + f"{config.__name__}.ConfigPatch[{id(self)}]", + default=(), + ) def __enter__(self) -> None: - if prior: - raise AssertionError( - "prior should be empty when entering ConfigPatch" - ) + prior: dict[str, Any] = {} for key in self.changes: # KeyError on invalid entry prior[key] = config.__getattr__(key) - for k, v in self.changes.items(): - config.__setattr__(k, v) + prior_stack = self._prior.get() + self._prior.set((*prior_stack, prior)) + try: + for k, v in self.changes.items(): + config.__setattr__(k, v) + except Exception: + self._prior.set(prior_stack) + raise def __exit__(self, exc_type, exc_val, exc_tb): # type: ignore[no-untyped-def] + prior_stack = self._prior.get() + if not prior_stack: + raise AssertionError( + "prior should not be empty when exiting ConfigPatch" + ) + prior = prior_stack[-1] + self._prior.set(prior_stack[:-1]) for k, v in prior.items(): config.__setattr__(k, v) - prior.clear() return ConfigPatch() @@ -778,13 +815,13 @@ def _make_closure_patcher(self, **changes: dict[str, Any]) -> Any: config = self._config def change() -> Callable[[], None]: - prior = {k: config[k].user_override for k in changes} + prior = {k: config[k].user_override.get() for k in changes} for k, v in changes.items(): - self._config[k].user_override = v + self._config[k].user_override.set(v) def revert() -> None: for k, v in prior.items(): - self._config[k].user_override = v + self._config[k].user_override.set(v) return revert @@ -853,15 +890,6 @@ def __delattr__(self, name: str) -> None: return self._config.__delattr__(self._prefix + name) -def patch_object(obj: object, name: str, value: object) -> object: - """ - Workaround `mock.patch.object` issue with ConfigModule - """ - if isinstance(obj, ConfigModule): - return obj.patch(name, value) - return mock.patch.object(obj, name, value) - - def get_tristate_env(name: str, default: Any = None) -> bool | None: value = os.environ.get(name) if value == "1": diff --git a/torch/utils/_inspect.py b/torch/utils/_inspect.py new file mode 100644 index 0000000000000..67d29b9bc7069 --- /dev/null +++ b/torch/utils/_inspect.py @@ -0,0 +1,90 @@ +import inspect +from typing import Any + + +def _signature_metadata( + sig: inspect.Signature, +) -> tuple[tuple[inspect.Parameter, ...], bool, int]: + """ + Returns tuple(sig.parameters.values()), if any has VAR_POSITIONAL or VAR_KEYWORD, and the max_positional + """ + params = tuple(sig.parameters.values()) + has_var_args = False + max_positional = 0 + + for p in params: + kind = p.kind + if kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + has_var_args = True + if kind in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ): + max_positional += 1 + + return params, has_var_args, max_positional + + +def _fast_bind( + sig: inspect.Signature, *args: Any, **kwargs: Any +) -> inspect.BoundArguments: + """ + Fast path for inspect.Signature.bind() for signatures without + VAR_POSITIONAL or VAR_KEYWORD parameters. Falls back to sig.bind() + for signatures that contain *args or **kwargs. + """ + params, has_var_args, max_positional = _signature_metadata(sig) + + # fallback for complex signatures + if has_var_args: + return sig.bind(*args, **kwargs) + + len_args = len(args) + + if len_args > max_positional: + raise TypeError( + f"Too many positional arguments: expected max {max_positional}, got {len_args}" + ) + + arguments: dict[str, Any] = {} + arg_i = 0 + + for p in params: + name = p.name + kind = p.kind + + if kind is inspect.Parameter.POSITIONAL_ONLY: + if name in kwargs: + raise TypeError( + f"Got some positional-only arguments passed as keyword arguments: '{name}'" + ) + if arg_i < len_args: + arguments[name] = args[arg_i] + arg_i += 1 + elif p.default is inspect.Parameter.empty: + raise TypeError(f"Missing required argument '{name}'") + + elif kind is inspect.Parameter.POSITIONAL_OR_KEYWORD: + if arg_i < len_args: + if name in kwargs: + raise TypeError(f"Multiple values for argument '{name}'") + arguments[name] = args[arg_i] + arg_i += 1 + elif name in kwargs: + arguments[name] = kwargs[name] + elif p.default is inspect.Parameter.empty: + raise TypeError(f"Missing required argument '{name}'") + + elif kind is inspect.Parameter.KEYWORD_ONLY: + if name in kwargs: + arguments[name] = kwargs[name] + elif p.default is inspect.Parameter.empty: + raise TypeError(f"Missing required argument '{name}'") + + # disallow extra keyword arguments not in the signature + # cause kwargs have been processed by sig.bind at the beginning + for name in kwargs: + if name not in sig.parameters: + raise TypeError(f"Got an unexpected keyword argument '{name}'") + + return inspect.BoundArguments(sig, arguments) # type: ignore[arg-type] diff --git a/torch/utils/_pytree.py b/torch/utils/_pytree.py index ee9af684bc6ce..de90d98f77951 100644 --- a/torch/utils/_pytree.py +++ b/torch/utils/_pytree.py @@ -22,6 +22,7 @@ import importlib import importlib.metadata import json +import logging import sys import threading import types @@ -47,6 +48,9 @@ from torch.torch_version import TorchVersion as _TorchVersion +log = logging.getLogger(__name__) + + if TYPE_CHECKING: import torch.utils._cxx_pytree as cxx_pytree @@ -617,10 +621,24 @@ def _private_register_pytree_node( from torch._library.opaque_object import is_opaque_type if isinstance(cls, type) and is_opaque_type(cls): - raise ValueError( - f"{cls} cannot be registered as a pytree as it has been " - "registered as an opaque object. Opaque objects must be pytree leaves." - ) + # TODO: remove this allowance once downstream callers stop calling + # register_constant on Enum subclasses. Enums are now natively + # supported as opaque value types and don't need pytree registration. + import enum + + if issubclass(cls, enum.Enum): + log.warning( + "%s is an Enum subclass and is now natively supported by " + "torch.compile as an opaque value type. Calling " + "register_constant() on Enum subclasses is deprecated and " + "will be an error in a future release.", + cls, + ) + else: + raise ValueError( + f"{cls} cannot be registered as a pytree as it has been " + "registered as an opaque object. Opaque objects must be pytree leaves." + ) with _NODE_REGISTRY_LOCK: if cls in SUPPORTED_NODES: @@ -1031,6 +1049,7 @@ def _deque_unflatten(values: Iterable[T], context: Context) -> deque[T]: STANDARD_DICT_TYPES: frozenset[type] = frozenset({dict, OrderedDict, defaultdict}) # pyrefly: ignore [no-matching-overload] BUILTIN_TYPES: frozenset[type] = frozenset( + # pyrefly: ignore [bad-argument-type] { tuple, list, diff --git a/torch/utils/_sympy/printers.py b/torch/utils/_sympy/printers.py index 963d0b5a20065..e9ca80812b04d 100644 --- a/torch/utils/_sympy/printers.py +++ b/torch/utils/_sympy/printers.py @@ -97,6 +97,9 @@ def _print_NegativeInfinity(self, expr: sympy.Expr) -> str: f"_print_NegativeInfinity not implemented for {type(self)}" ) + def _print_NaN(self, expr: sympy.Expr) -> str: + raise NotImplementedError(f"_print_NaN not implemented for {type(self)}") + def _print_FloorDiv(self, expr: sympy.Expr) -> str: raise NotImplementedError(f"_print_FloorDiv not implemented for {type(self)}") @@ -176,6 +179,9 @@ def _print_Infinity(self, expr: sympy.Expr) -> str: def _print_NegativeInfinity(self, expr: sympy.Expr) -> str: return "-math.inf" + def _print_NaN(self, expr: sympy.Expr) -> str: + return "math.nan" + # WARNING: this is dangerous for Triton, which has C-style modulus def _print_PythonMod(self, expr: sympy.Expr) -> str: return self.stringify(expr.args, " % ", PRECEDENCE["Atom"] - 0.5) @@ -191,8 +197,17 @@ def _print_IntTrueDiv(self, expr: sympy.Expr) -> str: return self.stringify(expr.args, " / ", PRECEDENCE["Atom"] - 0.5) def _helper_sqrt(self, expr: sympy.Expr) -> str: - # pyrefly: ignore [missing-attribute] - return f"math.sqrt({self._print(expr)})" + # NB: We use torch._sym_sqrt here instead of math.sqrt because the + # guard expression may be evaluated with SymInt/SymFloat inputs (e.g. + # during cache hit re-evaluation in evaluate_guards_expression). + # math.sqrt on a SymFloat triggers evaluate_expr which forces + # concretization/specialization of the symbol, creating spurious + # guards that didn't exist in the original program. + # torch._sym_sqrt properly propagates through the symbolic system + # without forcing specialization. + # See https://github.com/pytorch/pytorch/issues/152435 + # pyrefly: ignore [missing-attribute] + return f"torch._sym_sqrt({self._print(expr)})" def _print_OpaqueUnaryFn_sqrt(self, expr: sympy.Expr) -> str: return self._helper_sqrt(expr.args[0]) @@ -646,3 +661,6 @@ def _print_Infinity(self, expr: sympy.Expr) -> str: def _print_NegativeInfinity(self, expr: sympy.Expr) -> str: return f"-{self._print_Infinity(expr)}" + + def _print_NaN(self, expr: sympy.Expr) -> str: + return "std::numeric_limits::quiet_NaN()" diff --git a/torch/utils/_sympy/reference.py b/torch/utils/_sympy/reference.py index 3118e89ebfcf0..3879cef698154 100644 --- a/torch/utils/_sympy/reference.py +++ b/torch/utils/_sympy/reference.py @@ -177,7 +177,7 @@ def sqrt(x): @staticmethod def pow(a, b): - # pyrefly: ignore [bad-argument-type] + # pyrefly: ignore [bad-argument-count, bad-argument-type] return _keep_float(FloatPow)(a, b) @staticmethod @@ -340,6 +340,20 @@ def bitwise_or(a, b): def bitwise_xor(a, b): return a ^ b + @staticmethod + def expr_cond_pair(expr, cond): + return (expr, cond) + + @staticmethod + def piecewise(*pairs): + # Build nested sym_ite from right to left. + # Piecewise((e1, c1), (e2, c2), ..., (en, True)) becomes + # sym_ite(c1, e1, sym_ite(c2, e2, ... en)) + result = pairs[-1][0] + for expr, cond in reversed(pairs[:-1]): + result = torch.sym_ite(cond, expr, result) + return result + # Like PythonReferenceAnalysis, but some export-unfriendly choices of # operators to make things faster diff --git a/torch/utils/backend_registration.py b/torch/utils/backend_registration.py index 0e1eb82f8869e..f775cc0430d3c 100644 --- a/torch/utils/backend_registration.py +++ b/torch/utils/backend_registration.py @@ -186,9 +186,10 @@ def _generate_module_methods_for_privateuse1_backend(custom_backend_name: str) - ) def wrap_module_to( + # pyrefly: ignore [invalid-type-var] self: torch.nn.modules.module.T, device: int | torch.device | None = None, - ) -> torch.nn.modules.module.T: + ) -> torch.nn.modules.module.T: # pyrefly: ignore [invalid-type-var] r"""Move all model parameters and buffers to the custom device. This also makes associated parameters and buffers different objects. So diff --git a/torch/utils/checkpoint.py b/torch/utils/checkpoint.py index a21ffdb7fef29..d5f2bb039262b 100644 --- a/torch/utils/checkpoint.py +++ b/torch/utils/checkpoint.py @@ -2,6 +2,7 @@ from __future__ import annotations import contextlib +import itertools import platform import uuid import warnings @@ -138,10 +139,11 @@ def get_device_type() -> str: def _infer_device_type(*args): device_types = [] - def add_device_types(arg) -> None: + def add_device_types(arg): nonlocal device_types if isinstance(arg, torch.Tensor) and arg.device.type != "cpu": device_types.append(arg.device.type) + return arg tree_map(add_device_types, args) device_types_set = set(device_types) @@ -174,10 +176,11 @@ def get_device_states(*args) -> Tuple[List[int], List[torch.Tensor]]: # the conditionals short-circuit. fwd_device_ids = [] - def add_device_ids(arg) -> None: + def add_device_ids(arg): nonlocal fwd_device_ids if isinstance(arg, torch.Tensor) and arg.device.type not in {"cpu", "meta"}: fwd_device_ids.append(arg.get_device()) + return arg tree_map(add_device_ids, args) fwd_device_states = [] @@ -993,8 +996,7 @@ def _get_debug_context_and_cb() -> Tuple[Callable[[], Any], Callable[[Checkpoint # checkpointing mechanism. error_cb is invoked when an error is detected # during unpack. - # record_context_cpp is not support on non-linux non-x86_64 platforms - cpp_tb = platform.machine() == 'x86_64' and platform.system() == 'Linux' + cpp_tb = platform.machine() in ('x86_64', 'aarch64', 'arm64') and platform.system() == 'Linux' class CaptureLogs: def __init__(self) -> None: @@ -1069,6 +1071,10 @@ class _StopRecomputationError(Exception): class _recomputation_hook(torch.autograd.graph.saved_tensors_hooks): def __init__(self, target_frame_ref: ReferenceType, gid: GraphExecGroup | int) -> None: + # Dynamo guards on WeakKeyDictionary internals are unstable here + # (dict length/keys change every call), causing recompilation storms. + # with `.compile()` so we disable + @torch._dynamo.disable def pack_hook(x): x = x.detach() if x.requires_grad else x target_frame = target_frame_ref() @@ -1196,8 +1202,10 @@ def unpack_hook_with_error_cb(holder): def _is_compiling(func, args, kwargs): - # Check if we are under AOTAutograd tracing - # Checking that a functional mode is active should always do what we want + # Check if we are under AOTAutograd tracing or export tracing + # Checking that a proxy mode is active should always do what we want + if torch.compiler._is_non_strict_tracing(): + return False return torch._C._get_dispatch_mode(torch._C._TorchDispatchModeKey.PROXY) is not None @@ -1258,8 +1266,9 @@ class SelectiveCheckpointContext: >>> context_fn=context_fn, >>> ) """ - def __init__(self, *, is_recompute) -> None: + def __init__(self, *, is_recompute, op_output=None) -> None: self.is_recompute = is_recompute + self.op_output = op_output class CheckpointPolicy(enum.Enum): @@ -1316,25 +1325,29 @@ def ignore_compile_internals(cls): return True # Used together with _CachedTorchDispatchMode to implement SAC. - def __init__(self, policy_fn, storage) -> None: + def __init__(self, policy_fn, storage, ac_graph_id=None) -> None: self.policy_fn = policy_fn self.storage = storage + self.ac_graph_id = ac_graph_id def __torch_dispatch__(self, func, types, args=(), kwargs=None): - if func in SAC_IGNORED_OPS: - return func(*args, **kwargs) - kwargs = {} if kwargs is None else kwargs - policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=False), - func, *args, **kwargs) - if isinstance(policy, bool): - policy = _policy_from_bool(policy) - is_compiling = _is_compiling(func, args, kwargs) if is_compiling: - # Overwrite each node's "recompute" tag to add in the user annotation. - fx_traceback.current_meta["recompute"] = policy + fx_traceback.current_meta["ac_graph_id"] = self.ac_graph_id + fx_traceback.current_meta["recompute"] = CheckpointPolicy.PREFER_RECOMPUTE + + if func in SAC_IGNORED_OPS: + return func(*args, **kwargs) + + proxy_mode = None + graph_len_before = 0 + if is_compiling: + from torch.fx.experimental.proxy_tensor import get_proxy_mode + proxy_mode = get_proxy_mode() + if proxy_mode is not None: + graph_len_before = len(proxy_mode.tracer.graph.nodes) out = func(*args, **kwargs) @@ -1346,6 +1359,18 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None): else: any_ret_has_alias_info = any(ret.alias_info is not None for ret in func._schema.returns) + policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=False, op_output=out), + func, *args, **kwargs) + if isinstance(policy, bool): + policy = _policy_from_bool(policy) + + if is_compiling: + if proxy_mode is not None: + graph = proxy_mode.tracer.graph + num_new = len(graph.nodes) - graph_len_before + for node in itertools.islice(reversed(graph.nodes), num_new): + node.meta["recompute"] = policy + if policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE) or is_compiling: self.storage[func].append(tree_map(lambda x: _VersionWrapper(_maybe_detach(x, any_ret_has_alias_info)), out)) return out @@ -1572,6 +1597,8 @@ def _checkpoint_without_reentrant_generator( ), contextlib.nullcontext(), ) + error_on_nested_fx_trace = torch._dynamo.config.error_on_nested_fx_trace + is_non_strict_tracing = torch.compiler._is_non_strict_tracing() def recompute_fn(*args) -> None: # This will be called later during recomputation. This wrapping enables @@ -1590,7 +1617,20 @@ def recompute_fn(*args) -> None: device_autocast_ctx = torch.amp.autocast( device_type=device_type, **device_autocast_kwargs ) if torch.amp.is_autocast_available(device_type) else contextlib.nullcontext() - with device_autocast_ctx, torch.amp.autocast("cpu", **cpu_autocast_kwargs), recompute_context, device_ctx: # type: ignore[attr-defined] + nested_fx_trace_ctx = ( + torch._dynamo.config.patch( + error_on_nested_fx_trace=error_on_nested_fx_trace + ) + if is_non_strict_tracing + else contextlib.nullcontext() + ) + with ( + device_autocast_ctx, + torch.amp.autocast("cpu", **cpu_autocast_kwargs), + recompute_context, + device_ctx, + nested_fx_trace_ctx, + ): # type: ignore[attr-defined] fn(*args, **kwargs) new_frame = _CheckpointFrame( diff --git a/torch/utils/cpp_extension.py b/torch/utils/cpp_extension.py index a184eba0ce8a9..3e9d5d3823cbe 100644 --- a/torch/utils/cpp_extension.py +++ b/torch/utils/cpp_extension.py @@ -71,8 +71,8 @@ '12.5': ((6, 0, 0), (14, 0)), '12.6': ((6, 0, 0), (14, 0)), '12.7': ((6, 0, 0), (14, 0)), - '12.8': ((6, 0, 0), (14, 0)), - '12.9': ((6, 0, 0), (14, 0)), + '12.8': ((6, 0, 0), (15, 0)), + '12.9': ((6, 0, 0), (15, 0)), '13.0': ((6, 0, 0), (16, 0)), } @@ -2261,8 +2261,11 @@ def _jit_compile(name, hipified_sources = set() for source in sources: s_abs = os.path.abspath(source) - hipified_sources.add(hipify_result[s_abs].hipified_path if s_abs in hipify_result else s_abs) - + if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None: + hipified_s_abs = hipify_result[s_abs].hipified_path + else: + hipified_s_abs = s_abs + hipified_sources.add(hipified_s_abs) sources = list(hipified_sources) _write_ninja_file_and_build_library( diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py index 0b51561ca6e7f..4eaa755980787 100644 --- a/torch/utils/data/dataloader.py +++ b/torch/utils/data/dataloader.py @@ -423,8 +423,6 @@ def __init__( self.check_worker_number_rationality() - torch.set_vital("Dataloader", "enabled", "True") # type: ignore[attr-defined] - def _get_iterator(self) -> _BaseDataLoaderIter: if self.num_workers == 0: return _SingleProcessDataLoaderIter(self) @@ -606,31 +604,10 @@ def _create_warning_msg(num_worker_suggest, num_worker_created, cpuset_checked): return # try to compute a suggested max number of worker based on system's resource - max_num_worker_suggest = None - cpuset_checked = False - if hasattr(os, "sched_getaffinity"): - try: - max_num_worker_suggest = len(os.sched_getaffinity(0)) - cpuset_checked = True - except Exception: - pass - if max_num_worker_suggest is None: - # os.cpu_count() could return Optional[int] - # get cpu count first and check None in order to satisfy mypy check - cpu_count = os.cpu_count() - if cpu_count is not None: - max_num_worker_suggest = cpu_count - - if max_num_worker_suggest is None: - warnings.warn( - _create_warning_msg( - max_num_worker_suggest, self.num_workers, cpuset_checked - ), - stacklevel=2, - ) - return + max_num_worker_suggest = torch._utils.cpu_count() + cpuset_checked = hasattr(os, "sched_getaffinity") - if self.num_workers > max_num_worker_suggest: + if max_num_worker_suggest is None or self.num_workers > max_num_worker_suggest: warnings.warn( _create_warning_msg( max_num_worker_suggest, self.num_workers, cpuset_checked diff --git a/torch/utils/data/graph.py b/torch/utils/data/graph.py index 9820aac095ae3..3f8dbb17ceaea 100644 --- a/torch/utils/data/graph.py +++ b/torch/utils/data/graph.py @@ -11,6 +11,7 @@ __all__ = ["traverse", "traverse_dps"] DataPipe = IterDataPipe | MapDataPipe +# pyrefly: ignore [invalid-type-alias] DataPipeGraph = dict[int, tuple[DataPipe, "DataPipeGraph"]] diff --git a/torch/utils/flop_counter.py b/torch/utils/flop_counter.py index 76394f88cc32d..5f015a32f9c31 100644 --- a/torch/utils/flop_counter.py +++ b/torch/utils/flop_counter.py @@ -608,6 +608,85 @@ def _efficient_attention_backward_flop( ) +def _varlen_attn_forward_flop( + query, + key, + value, + cu_seq_q, + cu_seq_k, + max_q, + max_k, + *args, + out_val=None, + **kwargs, +) -> int: + """Count flops for varlen_attn forward.""" + sizes = _unpack_flash_attention_nested_shapes( + query=query, + key=key, + value=value, + cum_seq_q=cu_seq_q, + cum_seq_k=cu_seq_k if cu_seq_k is not None else cu_seq_q, + max_q=max_q, + max_k=max_k, + ) + return sum( + sdpa_flop_count(query_shape, key_shape, value_shape) + for query_shape, key_shape, value_shape, _ in sizes + ) + + +def _varlen_attn_out_flop( + out, + query, + key, + value, + cu_seq_q, + cu_seq_k, + max_q, + max_k, + *args, + out_val=None, + **kwargs, +) -> int: + """Count flops for varlen_attn_out forward.""" + return _varlen_attn_forward_flop( + query, key, value, cu_seq_q, cu_seq_k, max_q, max_k, + ) + + +def _varlen_attn_backward_flop( + grad_out, + query, + key, + value, + out, + lse, + cu_seq_q, + cu_seq_k, + max_q, + max_k, + *args, + out_val=None, + **kwargs, +) -> int: + """Count flops for varlen_attn backward.""" + sizes = _unpack_flash_attention_nested_shapes( + query=query, + key=key, + value=value, + grad_out=grad_out, + cum_seq_q=cu_seq_q, + cum_seq_k=cu_seq_k, + max_q=max_q, + max_k=max_k, + ) + return sum( + sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape) + for query_shape, key_shape, value_shape, grad_out_shape in sizes + ) + + flop_registry = { aten.mm: mm_flop, aten.addmm: addmm_flop, diff --git a/torch/utils/hipify/cuda_to_hip_mappings.py b/torch/utils/hipify/cuda_to_hip_mappings.py index b8bd0507e0652..9332d0a24f8ed 100644 --- a/torch/utils/hipify/cuda_to_hip_mappings.py +++ b/torch/utils/hipify/cuda_to_hip_mappings.py @@ -766,6 +766,7 @@ ("CU_MEM_HANDLE_TYPE_WIN32", "hipMemHandleTypeWin32"), ("CU_MEM_HANDLE_TYPE_WIN32_KMT", "hipMemHandleTypeWin32Kmt"), ("CU_MEM_LOCATION_TYPE_DEVICE", "hipMemLocationTypeDevice"), + ("CU_MEM_LOCATION_TYPE_HOST", "hipMemLocationTypeHost"), ("CU_MEM_LOCATION_TYPE_INVALID", "hipMemLocationTypeInvalid"), ("CU_MEM_OPERATION_TYPE_MAP", "hipMemOperationTypeMap"), ("CU_MEM_OPERATION_TYPE_UNMAP", "hipMemOperationTypeUnmap"), @@ -1360,6 +1361,10 @@ ("cudaEventBlockingSync", "hipEventBlockingSync"), ("cudaEventDisableTiming", "hipEventDisableTiming"), ("cudaEventInterprocess", "hipEventInterprocess"), + ("cudaEventRecordDefault", "hipEventRecordDefault"), + ("cudaEventRecordExternal", "hipEventRecordExternal"), + ("cudaEventWaitDefault", "hipEventWaitDefault"), + ("cudaEventWaitExternal", "hipEventWaitExternal"), ("cudaStreamCreate", "hipStreamCreate"), ("cudaStreamCreateWithFlags", "hipStreamCreateWithFlags"), ("cudaStreamCreateWithPriority", "hipStreamCreateWithPriority"), @@ -3392,6 +3397,23 @@ ("cudnnTensorDescriptor_t ", "miopenTensorDescriptor_t "), ("CUDNN_ENFORCE", "MIOPEN_ENFORCE"), ("CUDNN_CHECK", "MIOPEN_CHECK"), + # NVSHMEM → rocSHMEM mappings (only symbols used in hipified files: + # NVSHMEMSymmetricMemory.cpp and nvshmem_team_manager.hpp). + ("NVSHMEM_TEAM_INVALID", "rocshmem::ROCSHMEM_TEAM_INVALID"), + ("NVSHMEM_TEAM_WORLD", "rocshmem::ROCSHMEM_TEAM_WORLD"), + ("NVSHMEMX_INIT_WITH_UNIQUEID", "rocshmem::ROCSHMEM_INIT_WITH_UNIQUEID"), + + ("nvshmem_malloc", "rocshmem::rocshmem_malloc"), + ("nvshmem_free", "rocshmem::rocshmem_free"), + ("nvshmem_ptr", "rocshmem::rocshmem_ptr"), + ("nvshmem_team_t", "rocshmem::rocshmem_team_t"), + ("nvshmem_team_split_strided", "rocshmem::rocshmem_team_split_strided"), + + ("nvshmemx_uniqueid_t", "rocshmem::rocshmem_uniqueid_t"), + ("nvshmemx_get_uniqueid", "rocshmem::rocshmem_get_uniqueid"), + ("nvshmemx_init_attr", "rocshmem::rocshmem_init_attr"), + ("nvshmemx_init_attr_t", "rocshmem::rocshmem_init_attr_t"), + ("nvshmemx_set_attr_uniqueid_args", "rocshmem::rocshmem_set_attr_uniqueid_args"), ]) C10_MAPPINGS = collections.OrderedDict([ diff --git a/torch/utils/viz/MemoryViz.js b/torch/utils/viz/MemoryViz.js index c0ef85a65d2f3..4518c7755724d 100644 --- a/torch/utils/viz/MemoryViz.js +++ b/torch/utils/viz/MemoryViz.js @@ -1,3 +1,61 @@ +/** + * ================================================================================ + * MemoryViz.js - PyTorch Memory Visualization Tool + * ================================================================================ + * + * OVERVIEW: + * --------- + * This file contains the core visualization logic for PyTorch's memory profiler. + * It renders memory allocation timelines, stack traces, and provides interactive + * exploration of memory snapshots captured during model execution. + * + * KEY FEATURES: + * - Multiple visualization tabs/views for different memory analysis perspectives + * - Interactive stack trace display (supports both click and hover modes) + * - Zoom and brush controls for navigating large memory timelines + * - Support for loading memory snapshot files (.pickle format) + * + * ================================================================================ + * TESTING INSTRUCTIONS FOR ENGINEERS & AGENTS + * ================================================================================ + * + * 1. LOCAL TESTING SETUP: + * - Create a simple HTML file that references this JS file: + * + * + * + * MemoryViz Test + * + * + * + * + * + * - Serve locally using: python3 -m http.server 8888 + * - Open http://localhost:8888 in your browser + * + * 2. WHAT TO TEST: + * - Ensure ALL tabs/views render correctly and switch properly + * - Verify BOTH interaction modes work: + * * Click mode: stack traces appear on click + * * Hover mode: stack traces appear on mouseover + * - Test zoom and brush controls for timeline navigation + * - Verify memory allocation blocks are rendered and interactive + * + * 3. TEST DATA REQUIREMENTS: + * - DO NOT just test with small dummy .pickle files + * - Use realistic, decent-sized .pickle files (10-100+ MB range) + * - Large files stress-test rendering performance and memory handling + * - Test with snapshots from real model training/inference runs + * + * 4. COMMON ISSUES TO WATCH FOR: + * - Performance degradation with large snapshots + * - Stack trace popups not appearing or positioning incorrectly + * - Tab switching not updating the visualization properly + * - Zoom/brush state not persisting across interactions + * + * ================================================================================ + */ + 'use strict'; import * as d3 from "https://cdn.jsdelivr.net/npm/d3@7/+esm"; @@ -5,6 +63,9 @@ import {axisLeft} from "https://cdn.jsdelivr.net/npm/d3-axis@3/+esm"; import {scaleLinear} from "https://cdn.jsdelivr.net/npm/d3-scale@4/+esm"; import {zoom, zoomIdentity} from "https://cdn.jsdelivr.net/npm/d3-zoom@3/+esm"; import {brushX} from "https://cdn.jsdelivr.net/npm/d3-brush@3/+esm"; +import {process_alloc_data, isPrivatePoolId, formatSize, formatAddr, + elideRepeats, frameFilter, format_user_metadata, + format_forward_frames, format_frames} from "./process_alloc_data.js"; // Global configuration for trace interaction mode // 'hover' = show trace on hover (default) @@ -56,10 +117,6 @@ function Block(addr, size, requested_size, frames, free_requested, version, user return {addr, size, requested_size, frames, free_requested, version, user_metadata}; } -function isPrivatePoolId(pool_id) { - return pool_id && !(pool_id[0] === 0 && pool_id[1] === 0); -} - function EventSelector(outer, events, stack_info, memory_view) { const events_div = outer .append('div') @@ -119,25 +176,6 @@ function EventSelector(outer, events, stack_info, memory_view) { return es; } -function formatSize(num, showBytes = true) { - const orig = num; - // https://stackoverflow.com/questions/1094841/get-human-readable-version-of-file-size - const units = ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']; - for (const unit of units) { - if (Math.abs(num) < 1024.0) { - if (showBytes) { - return `${num.toFixed(1)}${unit}B (${orig} bytes)`; - } - return `${num.toFixed(1)}${unit}B`; - } - num /= 1024.0; - } - return `${num.toFixed(1)}YiB`; -} -function formatAddr(event) { - const prefix = event.action.startsWith('segment') ? 's\'' : 'b\''; - return `${prefix}${event.addr.toString(16)}_${event.version}`; -} function formatEvent(event) { const stream = event.stream === null ? '' : `\n (stream ${event.stream})`; @@ -808,580 +846,6 @@ function annotate_snapshot(snapshot) { } } -function elideRepeats(frames) { - const result = []; - const length = frames.length; - for (let i = 0; i < length; ) { - let j = i + 1; - const f = frames[i]; - while (j < length && f === frames[j]) { - j++; - } - switch (j - i) { - case 1: - result.push(f); - break; - case 2: - result.push(f, f); - break; - default: - result.push(f, ``); - break; - } - i = j; - } - return result; -} -function frameFilter({name, filename}) { - const omitFunctions = [ - 'unwind::unwind', - 'CapturedTraceback::gather', - 'gather_with_cpp', - '_start', - '__libc_start_main', - 'PyEval_', - 'PyObject_', - 'PyFunction_', - ]; - - const omitFilenames = [ - 'core/boxing', - '/Register', - '/Redispatch', - 'pythonrun.c', - 'Modules/main.c', - 'Objects/call.c', - 'Objects/methodobject.c', - 'pycore_ceval.h', - 'ceval.c', - 'cpython/abstract.h', - ]; - - for (const of of omitFunctions) { - if (name.includes(of)) { - return false; - } - } - - for (const of of omitFilenames) { - if (filename.includes(of)) { - return false; - } - } - - return true; -} - -function format_user_metadata(user_metadata) { - if (!user_metadata) { - return ''; - } - // Handle string metadata - if (typeof user_metadata === 'string') { - return `User Metadata:\n ${user_metadata}`; - } - // Handle object metadata - if (typeof user_metadata === 'object' && Object.keys(user_metadata).length === 0) { - return ''; - } - const metadata_lines = Object.entries(user_metadata) - .map(([key, value]) => ` ${key}: ${value}`); - return 'User Metadata:\n' + metadata_lines.join('\n'); -} - -function format_forward_frames(forward_frames) { - if (!forward_frames || forward_frames.length === 0) { - return ''; - } - // forward_frames is a list of strings (each string is a frame line from the forward pass) - // Each frame string already includes newlines, so we just join them directly - let frames_str = forward_frames.join(''); - // Ensure we don't have a trailing newline that could cause display issues - frames_str = frames_str.trimEnd(); - return `\n\n=== Forward Pass Stack Trace (where this tensor was created) ===\n${frames_str}`; -} - -function format_frames(frames) { - if (frames.length === 0) { - return ( - `This block has no frames. Potential causes:\n` + - `1) This block was allocated before _record_memory_history was enabled.\n` + - `2) The context or stacks passed to _record_memory_history does not include this block. Consider changing context to 'state', 'alloc', or 'all', or changing stacks to 'all'.\n` + - `3) This event occurred during backward, which has no python frames, and memory history did not include C++ frames. Use stacks='all' to record both C++ and python frames.` - ); - } - const frame_strings = frames - .filter(frameFilter) - .map(f => { - let frame_str = `${f.filename}:${f.line}:${f.name}`; - - // Add FX debug information if available - if (f.fx_node_op || f.fx_node_name || f.fx_node_target) { - const fx_parts = []; - if (f.fx_node_name) fx_parts.push(`node=${f.fx_node_name}`); - if (f.fx_node_op) fx_parts.push(`op=${f.fx_node_op}`); - if (f.fx_node_target) fx_parts.push(`target=${f.fx_node_target}`); - frame_str += `\n >> FX: ${fx_parts.join(', ')}`; - } - - if (f.fx_original_trace) { - frame_str += `\n >> Original Model Code:`; - const original_lines = f.fx_original_trace.trim().split('\n'); - // Show all lines of the original trace - for (const line of original_lines) { - frame_str += `\n ${line}`; - } - } - - return frame_str; - }); - return elideRepeats(frame_strings).join('\n'); -} - -function process_alloc_data(snapshot, device, plot_segments, max_entries, include_private_inactive = false) { - const elements = []; - const initially_allocated = []; - const actions = []; - const addr_to_alloc = {}; - - // Build sorted segment list for pool_id lookups by address - const device_segments = snapshot.segments - .filter(s => s.device === device) - .sort((a, b) => { - if (a.address === b.address) return 0; - return a.address < b.address ? -1 : 1; - }); - - // Binary search device_segments to find which segment contains addr, - // returning that segment's pool_id (or null if not found). - function find_pool_id(addr) { - let left = 0; - let right = device_segments.length - 1; - while (left <= right) { - const mid = Math.floor((left + right) / 2); - const seg = device_segments[mid]; - const seg_end = seg.address + (typeof seg.address === "bigint" ? BigInt(seg.total_size) : seg.total_size); - if (addr < seg.address) { - right = mid - 1; - } else if (addr >= seg_end) { - left = mid + 1; - } else { - return seg.segment_pool_id; - } - } - return null; - } - - const alloc = plot_segments ? 'segment_alloc' : 'alloc'; - const [free, free_completed] = plot_segments - ? ['segment_free', 'segment_free'] - : ['free', 'free_completed']; - for (const e of snapshot.device_traces[device]) { - switch (e.action) { - case alloc: - elements.push(e); - addr_to_alloc[e.addr] = elements.length - 1; - actions.push(elements.length - 1); - break; - case free: - case free_completed: - if (e.addr in addr_to_alloc) { - actions.push(addr_to_alloc[e.addr]); - delete addr_to_alloc[e.addr]; - } else { - elements.push(e); - initially_allocated.push(elements.length - 1); - actions.push(elements.length - 1); - } - break; - default: - break; - } - } - for (const seg of snapshot.segments) { - if (seg.device !== device) { - continue; - } - if (plot_segments) { - if (!(seg.address in addr_to_alloc)) { - const element = { - action: 'alloc', - addr: seg.address, - size: seg.total_size, - frames: [], - stream: seg.stream, - version: seg.version, - }; - elements.push(element); - initially_allocated.push(elements.length - 1); - } - } else { - const seg_is_private = isPrivatePoolId(seg.segment_pool_id); - for (const b of seg.blocks) { - const is_active = b.state === 'active_allocated'; - const is_private_inactive = include_private_inactive && seg_is_private && b.state === 'inactive'; - if ((is_active || is_private_inactive) && !(b.addr in addr_to_alloc)) { - const element = { - action: 'alloc', - addr: b.addr, - size: b.requested_size, - frames: b.frames, - stream: seg.stream, - version: b.version, - segment_pool_id: seg.segment_pool_id, - }; - elements.push(element); - initially_allocated.push(elements.length - 1); - } - } - } - } - // Resolve segment_pool_id for elements from trace events - for (const elem of elements) { - if (!elem.segment_pool_id) { - elem.segment_pool_id = find_pool_id(elem.addr); - } - } - - initially_allocated.reverse(); - // if there are no actions, the graph will be blank, - // but if there are existing allocations we do not want to hide them - // by having just one allocate action it will show a flat graph with all segments - if (actions.length === 0 && initially_allocated.length > 0) { - actions.push(initially_allocated.pop()); - } - - const current = []; - const current_data = []; - const data = []; - let max_size = 0; - - let total_mem = 0; - let total_summarized_mem = 0; - let timestep = 0; - - const max_at_time = []; - - const summarized_mem = { - elem: 'summarized', - timesteps: [], - offsets: [total_mem], - size: [], - color: 0, - }; - const summarized_elems = {}; - - function advance(n) { - summarized_mem.timesteps.push(timestep); - summarized_mem.offsets.push(total_mem); - summarized_mem.size.push(total_summarized_mem); - timestep += n; - for (let i = 0; i < n; i++) { - max_at_time.push(total_mem + total_summarized_mem); - } - } - - const sizes = elements - .map((x, i) => [x.size, i]) - .sort(([x, _xi], [y, _yi]) => y - x); - - const draw_elem = {}; - for (const [_s, e] of sizes.slice(0, max_entries)) { - draw_elem[e] = true; - } - - function add_allocation(elem) { - const element_obj = elements[elem]; - const size = element_obj.size; - current.push(elem); - let color = elem; - if (snapshot.categories.length > 0) { - color = snapshot.categories.indexOf(element_obj.category || 'unknown'); - } - const e = { - elem, - timesteps: [timestep], - offsets: [total_mem], - size, - color, - }; - current_data.push(e); - data.push(e); - total_mem += size; - element_obj.max_allocated_mem = total_mem + total_summarized_mem; - } - - // Pool state tracking (only active when include_private_inactive) - // Each pool is a single variable-size envelope in the global stack whose - // height = high-water mark. Colored stripes inside show active blocks. - const pools = {}; - const pool_active_elems = {}; - - function get_pool_key(elem_idx) { - const pid = elements[elem_idx].segment_pool_id; - if (!isPrivatePoolId(pid)) return null; - return `${pid[0]},${pid[1]}`; - } - - function get_or_create_pool(pool_key) { - if (!(pool_key in pools)) { - pools[pool_key] = { - max: 0, active: 0, envelope_data: null, - block_stack: [], // [{elem, size, inner_offset, stripe_data}] - }; - } - return pools[pool_key]; - } - - function elem_color(elem_idx) { - if (snapshot.categories.length > 0) { - return snapshot.categories.indexOf(elements[elem_idx].category || 'unknown'); - } - return elem_idx; - } - - // Shift all active stripes in a pool by delta - function shift_pool_stripes(pool, delta) { - for (const block of pool.block_stack) { - const s = block.stripe_data; - s.timesteps.push(timestep); - s.offsets.push(s.offsets.at(-1)); - s.timesteps.push(timestep + 3); - s.offsets.push(s.offsets.at(-1) + delta); - } - } - - // Shift global stack elements starting at idx by delta - function shift_elements_above(idx, delta) { - for (let j = idx; j < current.length; j++) { - const e = current_data[j]; - e.timesteps.push(timestep); - e.offsets.push(e.offsets.at(-1)); - e.timesteps.push(timestep + 3); - e.offsets.push(e.offsets.at(-1) + delta); - if (Array.isArray(e.size)) { - e.size.push(e.size.at(-1)); - e.size.push(e.size.at(-1)); - } - // If this is a pool envelope, also shift its stripes - const pk = typeof current[j] === 'string' && current[j].startsWith('pool:') - ? current[j].slice(5) : null; - if (pk && pk in pools) { - shift_pool_stripes(pools[pk], delta); - } - } - } - - // Grow the pool envelope from its current max to a new max - function grow_pool_envelope(pool, pool_key, new_active) { - if (new_active <= pool.max) return; - const delta = new_active - pool.max; - pool.max = new_active; - const env = pool.envelope_data; - env.timesteps.push(timestep); - env.offsets.push(env.offsets.at(-1)); - env.size.push(env.size.at(-1)); - env.timesteps.push(timestep + 3); - env.offsets.push(env.offsets.at(-1)); - env.size.push(pool.max); - const pidx = current.indexOf(`pool:${pool_key}`); - if (pidx >= 0) { - shift_elements_above(pidx + 1, delta); - } - total_mem += delta; - advance(3); - } - - for (const elem of initially_allocated) { - // Skip private pool blocks — they're handled by the pool envelope - if (include_private_inactive && get_pool_key(elem)) { - continue; - } - if (elem in draw_elem) { - add_allocation(elem); - } else { - total_summarized_mem += elements[elem].size; - summarized_elems[elem] = true; - } - } - - for (const elem of actions) { - const size = elements[elem].size; - const pool_key = include_private_inactive ? get_pool_key(elem) : null; - - if (pool_key) { - // Private pool element - if (!(elem in pool_active_elems)) { - // Pool alloc - pool_active_elems[elem] = pool_key; - const pool = get_or_create_pool(pool_key); - - // Create envelope on first alloc - if (pool.envelope_data === null) { - const env = { - elem: `pool:${pool_key}`, - timesteps: [timestep], - offsets: [total_mem], - size: [0], - color: 9, - }; - pool.envelope_data = env; - current.push(`pool:${pool_key}`); - current_data.push(env); - data.push(env); - } - - const inner_offset = pool.active; - pool.active += size; - - // Grow envelope first if needed (animates, shifts elements above) - if (pool.active > pool.max) { - grow_pool_envelope(pool, pool_key, pool.active); - } - - // Create stripe after growth so elements above have already shifted - const stripe = { - elem, - timesteps: [timestep], - offsets: [pool.envelope_data.offsets.at(-1) + inner_offset], - size, - color: elem_color(elem), - opacity: 0.5, - }; - pool.block_stack.push({elem, size, inner_offset, stripe_data: stripe}); - data.push(stripe); - advance(1); - elements[elem].max_allocated_mem = total_mem + total_summarized_mem; - } else { - // Pool free — envelope stays at max, end stripe, shift stripes above - const pool = pools[pool_key]; - const block_idx = pool.block_stack.findIndex(b => b.elem === elem); - if (block_idx >= 0) { - advance(1); - const block = pool.block_stack[block_idx]; - // End stripe - block.stripe_data.timesteps.push(timestep); - block.stripe_data.offsets.push(block.stripe_data.offsets.at(-1)); - - pool.block_stack.splice(block_idx, 1); - pool.active -= size; - - // Shift stripes above the freed block down within the pool - if (block_idx < pool.block_stack.length) { - for (let j = block_idx; j < pool.block_stack.length; j++) { - const b = pool.block_stack[j]; - b.inner_offset -= size; - const s = b.stripe_data; - s.timesteps.push(timestep); - s.offsets.push(s.offsets.at(-1)); - s.timesteps.push(timestep + 3); - s.offsets.push(pool.envelope_data.offsets.at(-1) + b.inner_offset); - } - advance(3); - } - } else { - pool.active -= size; - advance(1); - } - delete pool_active_elems[elem]; - } - max_size = Math.max(total_mem + total_summarized_mem, max_size); - continue; - } - - // Non-pool element (existing logic) - if (!(elem in draw_elem)) { - if (elem in summarized_elems) { - advance(1); - total_summarized_mem -= size; - summarized_elems[elem] = null; - } else { - total_summarized_mem += size; - summarized_elems[elem] = true; - advance(1); - } - continue; - } - const idx = current.findLastIndex(x => x === elem); - if (idx === -1) { - add_allocation(elem); - advance(1); - } else { - advance(1); - const removed = current_data[idx]; - removed.timesteps.push(timestep); - removed.offsets.push(removed.offsets.at(-1)); - current.splice(idx, 1); - current_data.splice(idx, 1); - - if (idx < current.length) { - shift_elements_above(idx, -size); - advance(3); - } - total_mem -= size; - } - max_size = Math.max(total_mem + total_summarized_mem, max_size); - } - - // Finalize: close all active elements - for (const elem of current_data) { - elem.timesteps.push(timestep); - elem.offsets.push(elem.offsets.at(-1)); - if (Array.isArray(elem.size)) { - elem.size.push(elem.size.at(-1)); - } - } - // Close active pool stripes - for (const pk in pools) { - for (const block of pools[pk].block_stack) { - const s = block.stripe_data; - s.timesteps.push(timestep); - s.offsets.push(s.offsets.at(-1)); - } - } - data.push(summarized_mem); - - return { - max_size, - allocations_over_time: data, - max_at_time, - summarized_mem, - elements_length: elements.length, - context_for_id: id => { - const elem = elements[id]; - let text = `Addr: ${formatAddr(elem)}`; - text = `${text}, Size: ${formatSize(elem.size)} allocation`; - text = `${text}, Total memory used after allocation: ${formatSize( - elem.max_allocated_mem, - )}`; - const context = elem?.compile_context ?? 'None'; - text = `${text}, Compile context: ${context}`; - if (elem.stream !== null) { - text = `${text}, stream ${elem.stream}`; - } - if (elem.segment_pool_id) { - text = `${text}, pool_id (${elem.segment_pool_id[0]}, ${elem.segment_pool_id[1]})`; - } else { - text = `${text}, pool_id unknown`; - } - if (elem.timestamp !== null) { - var d = new Date(elem.time_us / 1000); - text = `${text}, timestamp ${d}`; - } - if (!elem.action.includes('alloc')) { - text = `${text}\nalloc not recorded, stack trace for free:`; - } - const user_metadata_str = format_user_metadata(elem.user_metadata); - if (user_metadata_str) { - text = `${text}\n${user_metadata_str}`; - } - text = `${text}\n${format_frames(elem.frames)}`; - text = `${text}${format_forward_frames(elem.forward_frames)}`; - return text; - }, - }; -} - function MemoryPlot( svg, data, @@ -1503,11 +967,34 @@ function MemoryPlot( function ContextViewer(text, data) { let current_selected = null; + function restore_search_highlight(d) { + if (!d) return; + const addr = d.attr('data-search-match') === 'true'; + const frame = d.attr('data-frame-match') === 'true'; + if (addr && frame) { + d.attr('stroke', '#ff00ff') + .attr('stroke-width', 3) + .attr('stroke-dasharray', '6,3') + .attr('vector-effect', 'non-scaling-stroke'); + } else if (addr) { + d.attr('stroke', 'red') + .attr('stroke-width', 2) + .attr('stroke-dasharray', null) + .attr('vector-effect', 'non-scaling-stroke'); + } else if (frame) { + d.attr('stroke', '#2196F3') + .attr('stroke-width', 2) + .attr('stroke-dasharray', null) + .attr('vector-effect', 'non-scaling-stroke'); + } + } + return { default_selected: null, set_selected: d => { if (current_selected !== null) { - current_selected.attr('stroke', null).attr('stroke-width', null); + current_selected.attr('stroke', null).attr('stroke-width', null).attr('stroke-dasharray', null); + restore_search_highlight(current_selected); } if (d === null) { text.text(''); @@ -1649,6 +1136,22 @@ function create_trace_view( `Detail: ${max_entries} of ${data.elements_length} entries`, ); + d.append('span').text(' | '); + const search_input = d.append('input') + .attr('type', 'text') + .attr('placeholder', 'Search address (hex)...') + .attr('style', 'width: 180px; margin-left: 4px; font-family: monospace;'); + const search_label = d.append('label') + .attr('style', 'margin-left: 4px;'); + + d.append('span').text(' | '); + const frame_input = d.append('input') + .attr('type', 'text') + .attr('placeholder', 'Search stack frame...') + .attr('style', 'width: 200px; margin-left: 4px; font-family: monospace;'); + const frame_label = d.append('label') + .attr('style', 'margin-left: 4px;'); + const grid_container = dst .append('div') .attr( @@ -1685,6 +1188,61 @@ function create_trace_view( ); const delegate = ContextViewer(context_div.append('pre').text('none'), data); plot.set_delegate(delegate); + + function apply_search_highlights() { + const addr_query = search_input.node().value.toLowerCase().trim(); + const frame_query = frame_input.node().value.toLowerCase().trim(); + const polygons = plot_svg.selectAll('polygon'); + let addr_matches = 0; + let frame_matches = 0; + polygons.each(function () { + const dd = d3.select(this).datum(); + if (!dd || typeof dd.elem !== 'number') { + d3.select(this) + .attr('data-search-match', null) + .attr('data-frame-match', null); + return; + } + const ctx = data.context_for_id(dd.elem); + const ctx_lower = ctx.toLowerCase(); + const addr_hit = addr_query && ctx_lower.includes(addr_query); + const frame_hit = frame_query && ctx_lower.includes(frame_query); + d3.select(this) + .attr('data-search-match', addr_hit ? 'true' : null) + .attr('data-frame-match', frame_hit ? 'true' : null); + if (addr_hit && frame_hit) { + d3.select(this) + .attr('stroke', '#ff00ff') + .attr('stroke-width', 3) + .attr('stroke-dasharray', '6,3') + .attr('vector-effect', 'non-scaling-stroke'); + } else if (addr_hit) { + d3.select(this) + .attr('stroke', 'red') + .attr('stroke-width', 2) + .attr('stroke-dasharray', null) + .attr('vector-effect', 'non-scaling-stroke'); + } else if (frame_hit) { + d3.select(this) + .attr('stroke', '#2196F3') + .attr('stroke-width', 2) + .attr('stroke-dasharray', null) + .attr('vector-effect', 'non-scaling-stroke'); + } else { + d3.select(this) + .attr('stroke', null) + .attr('stroke-width', null) + .attr('stroke-dasharray', null); + } + if (addr_hit) addr_matches++; + if (frame_hit) frame_matches++; + }); + search_label.text(addr_query ? `${addr_matches} match${addr_matches !== 1 ? 'es' : ''}` : ''); + frame_label.text(frame_query ? `${frame_matches} match${frame_matches !== 1 ? 'es' : ''}` : ''); + } + + search_input.on('input', apply_search_highlights); + frame_input.on('input', apply_search_highlights); } function create_settings_view(dst, snapshot, device) { diff --git a/torch/utils/viz/process_alloc_data.js b/torch/utils/viz/process_alloc_data.js new file mode 100644 index 0000000000000..6e5feda1046ea --- /dev/null +++ b/torch/utils/viz/process_alloc_data.js @@ -0,0 +1,902 @@ +// Pure data-processing functions for PyTorch memory visualization. +// Extracted from MemoryViz.js so they can be tested independently (no d3/DOM deps). +// +// This file is the single source of truth for these functions: +// - MemoryViz.js imports them via ESM: import {...} from "./process_alloc_data.js" +// - Node.js tests load this file by stripping the export line and eval-ing +// +// TRACE EVENT ACTIONS (from c10/core/CachingDeviceAllocator.h TraceEntry::Action): +// +// "alloc" - Sub-allocation returned to user from the caching allocator. +// Recorded in alloc_found_block() (CUDACachingAllocator.cpp:1834). +// +// "free_requested" - User code called free (tensor out of scope). The block may not +// be immediately returned to the free pool if it's in use on +// another stream via record_stream. +// Recorded in free() (CUDACachingAllocator.cpp:2123). +// +// "free_completed" - Block actually returned to the allocator's free pool. For simple +// cases this fires immediately after free_requested. For cross-stream +// blocks, deferred until CUDA events confirm all streams are done. +// Recorded in free_block() (CUDACachingAllocator.cpp:3148). +// +// "segment_alloc" - New segment allocated from OS via cudaMalloc (or cuMemCreate for +// expandable segments). +// Recorded in alloc_from_expandable_segment() (CUDACachingAllocator.cpp:3548). +// +// "segment_free" - Segment returned to OS via cudaFree. Happens during empty_cache() +// or defragmentation. Only for non-expandable segments. +// Recorded in release_block() (CUDACachingAllocator.cpp:3686). +// +// "segment_map" - Physical pages mapped into an expandable segment via cuMemMap. +// The segment grows. Only with expandable segments enabled. +// Recorded in alloc_from_expandable_segment() (CUDACachingAllocator.cpp:3092). +// +// "segment_unmap" - Physical pages unmapped from an expandable segment via cuMemUnmap. +// Virtual address range retained, physical memory returned to OS. +// Only with expandable segments. Causes "pool_id unknown" for any +// trace events whose addresses fall in the unmapped range, since +// the segment no longer exists at snapshot time. +// Recorded in unmap_block() (CUDACachingAllocator.cpp:3790). +// +// "snapshot" - A call to torch.cuda.memory._snapshot(). Timestamp marker to +// correlate trace events with snapshot state. addr=0. +// Recorded in snapshot() (CUDACachingAllocator.cpp:2689). +// +// "oom" - Allocator failed to satisfy an allocation after all retries. +// addr=device_free (bytes free on GPU), size=requested allocation. +// Recorded in malloc() (CUDACachingAllocator.cpp:1629). +// +// HOW SEGMENT EVENTS ARE USED IN VISUALIZATION: +// +// The snapshot pickle contains two separate data sources: +// 1. device_traces - Ring buffer of TraceEntry actions (alloc, free, segment_map, etc.) +// 2. segments - Point-in-time dump of all segments/blocks at _snapshot() time +// +// Block-level views ("Active Memory Timeline", "Allocated Memory (incl. Private Pools)"): +// - process_alloc_data matches "alloc" and "free_completed" from device_traces. +// - segment_alloc/segment_free/segment_map/segment_unmap are ignored (skipped in switch). +// - The segments snapshot is used only to resolve pool_id via find_pool_id(). +// +// Segment-level view ("Active Cached Segment Timeline"): +// - process_alloc_data is called with plot_segments=true. +// - Matches "segment_alloc" and "segment_free" instead of alloc/free. +// - segment_map/segment_unmap are NOT matched (they don't appear in the switch). +// - Segments from the snapshot that weren't seen in the trace are added as +// initially_allocated (Phase 2). +// +// Allocator State History ("Allocator State History"): +// - EventSelector lists ALL trace events including segment_map/segment_unmap. +// - MemoryView renders the segment/block layout from the segments snapshot. +// - Clicking an event in the list redraws the layout at that point in time. +// +// Ring buffer overflow: +// - All trace event types share the same ring buffer. When it overflows, older +// events are overwritten. The allocator_settings.trace_alloc_overflowed flag +// indicates this happened, and trace_alloc_max_entries gives the buffer size. +// - Segment snapshot data (segments array) is NOT affected by ring buffer overflow. +// - The segment snapshot is always complete regardless of overflow. + +/** + * Returns true if pool_id represents a private (user-created) memory pool, + * as opposed to the default pool [0, 0]. + * + * @param {number[]|null} pool_id - Two-element array [owner_id, pool_id] from + * the CUDA caching allocator. The default pool is [0, 0]; any other non-null + * value is a private pool (e.g. FSDP's MemPool). + * @returns {boolean} + */ +function isPrivatePoolId(pool_id) { + return pool_id && !(pool_id[0] === 0 && pool_id[1] === 0); +} + +/** + * Formats a byte count as a human-readable string (e.g. "1.5GiB (1610612736 bytes)"). + * + * @param {number} num - Size in bytes. + * @param {boolean} [showBytes=true] - Whether to include the raw byte count in parentheses. + * @returns {string} + */ +function formatSize(num, showBytes = true) { + const orig = num; + // https://stackoverflow.com/questions/1094841/get-human-readable-version-of-file-size + const units = ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']; + for (const unit of units) { + if (Math.abs(num) < 1024.0) { + if (showBytes) { + return `${num.toFixed(1)}${unit}B (${orig} bytes)`; + } + return `${num.toFixed(1)}${unit}B`; + } + num /= 1024.0; + } + return `${num.toFixed(1)}YiB`; +} + +/** + * Formats a trace event's address as a display string like "b'7f4c00000_3". + * Segment-level events get an "s'" prefix, block-level events get "b'". + * + * @param {{action: string, addr: number|BigInt, version: number}} event + * @returns {string} + */ +function formatAddr(event) { + const prefix = event.action.startsWith('segment') ? 's\'' : 'b\''; + return `${prefix}${event.addr.toString(16)}_${event.version}`; +} + +/** + * Collapses consecutive duplicate strings in an array. If a string appears + * N > 2 times in a row, it's replaced with [str, ""]. + * Used to compress repetitive stack frames in the display. + * + * @param {string[]} frames + * @returns {string[]} + */ +function elideRepeats(frames) { + const result = []; + const length = frames.length; + for (let i = 0; i < length; ) { + let j = i + 1; + const f = frames[i]; + while (j < length && f === frames[j]) { + j++; + } + switch (j - i) { + case 1: + result.push(f); + break; + case 2: + result.push(f, f); + break; + default: + result.push(f, ``); + break; + } + i = j; + } + return result; +} + +/** + * Returns false for stack frames that are internal runtime noise + * (e.g. Python interpreter internals, C++ dispatch machinery). + * Used as a filter predicate on frame arrays. + * + * @param {{name: string, filename: string}} frame + * @returns {boolean} + */ +function frameFilter({name, filename}) { + const omitFunctions = [ + 'unwind::unwind', + 'CapturedTraceback::gather', + 'gather_with_cpp', + '_start', + '__libc_start_main', + 'PyEval_', + 'PyObject_', + 'PyFunction_', + ]; + + const omitFilenames = [ + 'core/boxing', + '/Register', + '/Redispatch', + 'pythonrun.c', + 'Modules/main.c', + 'Objects/call.c', + 'Objects/methodobject.c', + 'pycore_ceval.h', + 'ceval.c', + 'cpython/abstract.h', + ]; + + for (const of of omitFunctions) { + if (name.includes(of)) { + return false; + } + } + + for (const of of omitFilenames) { + if (filename.includes(of)) { + return false; + } + } + + return true; +} + +/** + * Formats user-attached metadata (from torch.cuda.memory._record_memory_history) + * as a display string. Returns '' if no metadata is present. + * + * @param {string|Object|null|undefined} user_metadata + * @returns {string} + */ +function format_user_metadata(user_metadata) { + if (!user_metadata) { + return ''; + } + if (typeof user_metadata === 'string') { + return `User Metadata:\n ${user_metadata}`; + } + if (typeof user_metadata === 'object' && Object.keys(user_metadata).length === 0) { + return ''; + } + const metadata_lines = Object.entries(user_metadata) + .map(([key, value]) => ` ${key}: ${value}`); + return 'User Metadata:\n' + metadata_lines.join('\n'); +} + +/** + * Formats the forward-pass stack trace (captured via torch.autograd) as a + * display string showing where a tensor was originally created. + * + * @param {string[]|null|undefined} forward_frames + * @returns {string} + */ +function format_forward_frames(forward_frames) { + if (!forward_frames || forward_frames.length === 0) { + return ''; + } + let frames_str = forward_frames.join(''); + frames_str = frames_str.trimEnd(); + return `\n\n=== Forward Pass Stack Trace (where this tensor was created) ===\n${frames_str}`; +} + +/** + * Formats an array of stack frames into a human-readable string. + * Filters out runtime noise via frameFilter, annotates FX graph debug info + * when available, and collapses consecutive duplicate frames. + * + * @param {{filename: string, line: number, name: string, + * fx_node_op?: string, fx_node_name?: string, + * fx_node_target?: string, fx_original_trace?: string}[]} frames + * @returns {string} + */ +function format_frames(frames) { + if (frames.length === 0) { + return ( + `This block has no frames. Potential causes:\n` + + `1) This block was allocated before _record_memory_history was enabled.\n` + + `2) The context or stacks passed to _record_memory_history does not include this block. Consider changing context to 'state', 'alloc', or 'all', or changing stacks to 'all'.\n` + + `3) This event occurred during backward, which has no python frames, and memory history did not include C++ frames. Use stacks='all' to record both C++ and python frames.\n` + + `4) This block was reconstructed from the allocator's segment snapshot (not from a trace event). The snapshot records which blocks exist at the moment _snapshot() is called, but does not carry stack frames. This typically happens for blocks that were allocated before tracing started and never freed, or for inactive blocks in private memory pools.\n` + + `5) The original alloc event was evicted from the trace ring buffer (older entries are overwritten when the buffer is full). Increase the max_entries argument to _record_memory_history to retain more events.` + ); + } + const frame_strings = frames + .filter(frameFilter) + .map(f => { + let frame_str = `${f.filename}:${f.line}:${f.name}`; + + if (f.fx_node_op || f.fx_node_name || f.fx_node_target) { + const fx_parts = []; + if (f.fx_node_name) fx_parts.push(`node=${f.fx_node_name}`); + if (f.fx_node_op) fx_parts.push(`op=${f.fx_node_op}`); + if (f.fx_node_target) fx_parts.push(`target=${f.fx_node_target}`); + frame_str += `\n >> FX: ${fx_parts.join(', ')}`; + } + + if (f.fx_original_trace) { + frame_str += `\n >> Original Model Code:`; + const original_lines = f.fx_original_trace.trim().split('\n'); + for (const line of original_lines) { + frame_str += `\n ${line}`; + } + } + + return frame_str; + }); + return elideRepeats(frame_strings).join('\n'); +} + +/** + * Transforms a memory snapshot into a stacked-area timeline suitable for + * rendering by MemoryPlot. This is the core data-processing function behind + * the "Active Memory Timeline" and "Allocated Memory (incl. Private Pools)" + * visualization tabs. + * + * HIGH-LEVEL ALGORITHM: + * + * 1. TRACE EVENT MATCHING: Scans device_traces to pair alloc events with their + * corresponding free_completed events (by address). Events whose matching + * alloc was lost (e.g. ring buffer wrap) become "initially_allocated" — + * blocks assumed to exist at the start of the trace. + * + * 2. SEGMENT SNAPSHOT: Supplements trace data with the current segment state. + * Blocks marked active_allocated (or inactive in private pools, when + * include_private_inactive=true) that weren't seen in the trace are also + * added as initially_allocated. + * + * 3. DETAIL LIMITING: Only the largest max_entries elements get individual + * rectangles in the plot. Smaller elements are aggregated into a single + * "summarized" band to keep rendering fast. + * + * 4. STACKED AREA CONSTRUCTION: Replays alloc/free events in order, building + * a stacked-area dataset where each element has timesteps, y-offsets, and + * a size. Elements are stacked bottom-to-top; frees remove from the stack + * and shift elements above downward. + * + * 5. PRIVATE POOL ENVELOPES (include_private_inactive=true): Each private pool + * (e.g. FSDP's MemPool) gets a single gray "envelope" rectangle whose + * height is the pool's high-water mark. Active blocks within the pool are + * rendered as colored stripes inside the envelope. The envelope only grows + * (never shrinks), representing reserved capacity. + * + * Initially-allocated private pool blocks are PRE-LOADED into pool state + * so that when their free event appears in the trace, they are correctly + * recognized as frees (not misinterpreted as new allocations). + * + * NOTE ON FREE EVENT MATCHING: The C++ allocator emits 'free_requested' and + * 'free_completed' for each deallocation. This function matches against 'free' + * (which no longer appears in modern traces — effectively dead code) and + * 'free_completed'. Only free_completed does the actual matching. Matching + * both 'free_requested' AND 'free_completed' would cause double-processing + * since they share the same address. + * + * @param {Object} snapshot - Memory snapshot from torch.cuda.memory._snapshot(). + * @param {Object[]} snapshot.segments - Current allocator segment state. + * @param {Object[][]} snapshot.device_traces - Per-device arrays of trace events. + * Each event has {action, addr, size, frames, stream, segment_pool_id?, ...}. + * @param {string[]} snapshot.categories - Category names for color-coding. + * @param {number} device - Device index into snapshot.device_traces. + * @param {boolean} plot_segments - If true, plot segment-level (cudaMalloc) + * events instead of sub-allocation events. + * @param {number} max_entries - Maximum number of elements to render individually. + * Elements beyond this limit are aggregated into the "summarized" band. + * @param {boolean} [include_private_inactive=false] - If true, include inactive + * blocks from private pools and render pool envelopes. Used by the + * "Allocated Memory (incl. Private Pools)" tab. + * + * @returns {{ + * max_size: number, + * allocations_over_time: Object[], + * max_at_time: number[], + * summarized_mem: Object, + * elements_length: number, + * context_for_id: function(number): string + * }} + * - max_size: peak total memory observed during the action replay (used for + * y-axis scaling). Note: this is only updated inside the action loop, so + * the initial state from initially_allocated may not be reflected here + * (use max_at_time for the true peak). + * - allocations_over_time: array of stacked-area data objects, each with + * {elem, timesteps[], offsets[], size, color}. + * - max_at_time: total memory at each timestep (for minimap rendering). + * - summarized_mem: the aggregated band for small elements. + * - elements_length: total number of unique allocation elements. + * - context_for_id: function that returns a human-readable description + * string for a given element index (address, size, stack trace, etc.). + */ +function process_alloc_data(snapshot, device, plot_segments, max_entries, include_private_inactive = false) { + const elements = []; + // Contains two types of blocks + // 1. free without alloc in trace + // 2. actively allocated in segments, but no matching alloc in trace + const initially_allocated = []; + const actions = []; + const addr_to_alloc = {}; + + const device_segments = snapshot.segments + .filter(s => s.device === device) + .sort((a, b) => { + if (a.address === b.address) return 0; + return a.address < b.address ? -1 : 1; + }); + + // Binary search to find which segment contains a given address. + function find_pool_id(addr) { + let left = 0; + let right = device_segments.length - 1; + while (left <= right) { + const mid = Math.floor((left + right) / 2); + const seg = device_segments[mid]; + const seg_end = seg.address + (typeof seg.address === "bigint" ? BigInt(seg.total_size) : seg.total_size); + if (addr < seg.address) { + right = mid - 1; + } else if (addr >= seg_end) { + left = mid + 1; + } else { + return seg.segment_pool_id; + } + } + return null; + } + + const alloc = plot_segments ? 'segment_alloc' : 'alloc'; + const [free, free_completed] = plot_segments + ? ['segment_free', 'segment_free'] + : ['free', 'free_completed']; + for (const e of snapshot.device_traces[device]) { + switch (e.action) { + case alloc: + elements.push(e); + addr_to_alloc[e.addr] = elements.length - 1; + actions.push(elements.length - 1); + break; + case free: + case free_completed: + if (e.addr in addr_to_alloc) { + // Matched: reuse the element from the alloc event + actions.push(addr_to_alloc[e.addr]); + delete addr_to_alloc[e.addr]; + } else { + // Unmatched free: alloc happened before recording (or was evicted + // from the ring buffer). Create a new element from the free event; + // its stack trace will show the free site, not the alloc site. + elements.push(e); + initially_allocated.push(elements.length - 1); + actions.push(elements.length - 1); + } + break; + default: + break; + } + } + + // --- Phase 2: Add elements from the snapshot --- + for (const seg of snapshot.segments) { + if (seg.device !== device) { + continue; + } + if (plot_segments) { + if (!(seg.address in addr_to_alloc)) { + const element = { + action: 'alloc', + addr: seg.address, + size: seg.total_size, + frames: [], + stream: seg.stream, + version: seg.version, + }; + elements.push(element); + initially_allocated.push(elements.length - 1); + } + } else { + for (const b of seg.blocks) { + const addr = b.addr ?? b.address; + if (b.state === 'active_allocated' && !(addr in addr_to_alloc)) { + const element = { + action: 'alloc', + addr, + size: b.requested_size, + frames: b.frames, + stream: seg.stream, + version: b.version, + segment_pool_id: seg.segment_pool_id, + ghost: true, + }; + elements.push(element); + initially_allocated.push(elements.length - 1); + } + } + } + } + + // Resolve pool IDs for trace elements by looking up which segment they fall in + for (const elem of elements) { + if (!elem.segment_pool_id) { + elem.segment_pool_id = find_pool_id(elem.addr); + } + } + + initially_allocated.reverse(); + // If there are no trace actions but there are existing allocations, + // show a flat graph with the initial state + if (actions.length === 0 && initially_allocated.length > 0) { + actions.push(initially_allocated.pop()); + } + + // --- Phase 3: Build the stacked-area timeline --- + const current = []; // stack of element indices (bottom to top) + const current_data = []; // parallel array of visualization data objects + const data = []; // all data objects (including completed ones) + let max_size = 0; + + let total_mem = 0; + let total_summarized_mem = 0; + let timestep = 0; + + const max_at_time = []; + + const summarized_mem = { + elem: 'summarized', + timesteps: [], + offsets: [total_mem], + size: [], + color: 0, + }; + const summarized_elems = {}; + + // Record the current memory state and advance time by n steps + function advance(n) { + summarized_mem.timesteps.push(timestep); + summarized_mem.offsets.push(total_mem); + summarized_mem.size.push(total_summarized_mem); + timestep += n; + for (let i = 0; i < n; i++) { + max_at_time.push(total_mem + total_summarized_mem); + } + } + + // Only render the largest max_entries elements individually; + // everything else goes into the summarized band + const sizes = elements + .map((x, i) => [x.size, i]) + .sort(([x, _xi], [y, _yi]) => y - x); + + const draw_elem = {}; + for (const [_s, e] of sizes.slice(0, max_entries)) { + draw_elem[e] = true; + } + + // Push an element onto the memory stack + function add_allocation(elem) { + const element_obj = elements[elem]; + const size = element_obj.size; + current.push(elem); + let color = elem; + if (snapshot.categories.length > 0) { + color = snapshot.categories.indexOf(element_obj.category || 'unknown'); + } + const e = { + elem, + timesteps: [timestep], + offsets: [total_mem], + size, + color, + }; + if (element_obj.ghost) e.ghost = true; + current_data.push(e); + data.push(e); + total_mem += size; + element_obj.max_allocated_mem = total_mem + total_summarized_mem; + } + + // --- Pool envelope tracking (only when include_private_inactive=true) --- + // Each private pool gets a gray envelope whose height = high-water mark. + // Active blocks are rendered as colored stripes within the envelope. + const pools = {}; + const pool_active_elems = {}; + + function get_pool_key(elem_idx) { + const pid = elements[elem_idx].segment_pool_id; + if (!isPrivatePoolId(pid)) return null; + const stream = elements[elem_idx].stream; + return `${pid[0]},${pid[1]},s${stream}`; + } + + function get_or_create_pool(pool_key) { + if (!(pool_key in pools)) { + pools[pool_key] = { + max: 0, active: 0, envelope_data: null, + block_stack: [], // [{elem, size, inner_offset, stripe_data}] + }; + } + return pools[pool_key]; + } + + function elem_color(elem_idx) { + if (snapshot.categories.length > 0) { + return snapshot.categories.indexOf(elements[elem_idx].category || 'unknown'); + } + return elem_idx; + } + + function shift_pool_stripes(pool, delta) { + for (const block of pool.block_stack) { + const s = block.stripe_data; + s.timesteps.push(timestep); + s.offsets.push(s.offsets.at(-1)); + s.timesteps.push(timestep + 3); + s.offsets.push(s.offsets.at(-1) + delta); + } + } + + // Animate shifting all elements above idx by delta (used when an element + // is inserted or removed from the middle of the stack) + function shift_elements_above(idx, delta) { + for (let j = idx; j < current.length; j++) { + const e = current_data[j]; + e.timesteps.push(timestep); + e.offsets.push(e.offsets.at(-1)); + e.timesteps.push(timestep + 3); + e.offsets.push(e.offsets.at(-1) + delta); + if (Array.isArray(e.size)) { + e.size.push(e.size.at(-1)); + e.size.push(e.size.at(-1)); + } + const pk = typeof current[j] === 'string' && current[j].startsWith('pool:') + ? current[j].slice(5) : null; + if (pk && pk in pools) { + shift_pool_stripes(pools[pk], delta); + } + } + } + + // Grow a pool envelope to accommodate new_active bytes. + // The envelope only grows (never shrinks) — it represents reserved capacity. + function grow_pool_envelope(pool, pool_key, new_active) { + if (new_active <= pool.max) return; + const delta = new_active - pool.max; + pool.max = new_active; + const env = pool.envelope_data; + env.timesteps.push(timestep); + env.offsets.push(env.offsets.at(-1)); + env.size.push(env.size.at(-1)); + env.timesteps.push(timestep + 3); + env.offsets.push(env.offsets.at(-1)); + env.size.push(pool.max); + const pidx = current.indexOf(`pool:${pool_key}`); + if (pidx >= 0) { + shift_elements_above(pidx + 1, delta); + } + total_mem += delta; + advance(3); + } + + // --- Process initially_allocated elements --- + // Private pool blocks are pre-loaded into pool state at timestep 0 (no + // animation) so that their envelope starts at the correct initial size + // and free events are correctly recognized as frees. + for (const elem of initially_allocated) { + if (include_private_inactive && get_pool_key(elem)) { + const pk = get_pool_key(elem); + const size = elements[elem].size; + const pool = get_or_create_pool(pk); + pool_active_elems[elem] = pk; + + if (pool.envelope_data === null) { + const env = { + elem: `pool:${pk}`, + timesteps: [0], + offsets: [total_mem], + size: [0], + color: 9, + }; + pool.envelope_data = env; + current.push(`pool:${pk}`); + current_data.push(env); + data.push(env); + } + + const inner_offset = pool.active; + pool.active += size; + + // Grow envelope directly without animation — these blocks pre-exist + if (pool.active > pool.max) { + const delta = pool.active - pool.max; + pool.max = pool.active; + const env = pool.envelope_data; + env.size[env.size.length - 1] = pool.max; + total_mem += delta; + const pidx = current.indexOf(`pool:${pk}`); + if (pidx >= 0) { + for (let j = pidx + 1; j < current.length; j++) { + const e = current_data[j]; + e.offsets[e.offsets.length - 1] += delta; + } + } + } + + const stripe = { + elem, + timesteps: [0], + offsets: [pool.envelope_data.offsets.at(-1) + inner_offset], + size, + color: elem_color(elem), + opacity: 0.5, + ghost: elements[elem].ghost || false, + }; + pool.block_stack.push({elem, size, inner_offset, stripe_data: stripe}); + data.push(stripe); + continue; + } + if (elem in draw_elem) { + add_allocation(elem); + } else { + total_summarized_mem += elements[elem].size; + summarized_elems[elem] = true; + } + } + + // Fix up pool stripe offsets — stripes are not in current_data so they + // don't get shifted when other pools grow during initially_allocated + // processing. Recompute from the envelope's final offset. + for (const pk in pools) { + const p = pools[pk]; + if (!p.envelope_data) continue; + const env_offset = p.envelope_data.offsets.at(-1); + for (const block of p.block_stack) { + const s = block.stripe_data; + for (let i = 0; i < s.offsets.length; i++) { + s.offsets[i] = env_offset + block.inner_offset; + } + } + } + + // --- Replay alloc/free actions to build the timeline --- + for (const elem of actions) { + const size = elements[elem].size; + const pool_key = include_private_inactive ? get_pool_key(elem) : null; + + if (pool_key) { + // --- Private pool element --- + if (!(elem in pool_active_elems)) { + // Pool alloc: add to pool, grow envelope if needed + pool_active_elems[elem] = pool_key; + const pool = get_or_create_pool(pool_key); + + if (pool.envelope_data === null) { + const env = { + elem: `pool:${pool_key}`, + timesteps: [timestep], + offsets: [total_mem], + size: [0], + color: 9, + }; + pool.envelope_data = env; + current.push(`pool:${pool_key}`); + current_data.push(env); + data.push(env); + } + + const inner_offset = pool.active; + pool.active += size; + + if (pool.active > pool.max) { + grow_pool_envelope(pool, pool_key, pool.active); + } + + const stripe = { + elem, + timesteps: [timestep], + offsets: [pool.envelope_data.offsets.at(-1) + inner_offset], + size, + color: elem_color(elem), + opacity: 0.5, + }; + pool.block_stack.push({elem, size, inner_offset, stripe_data: stripe}); + data.push(stripe); + advance(1); + elements[elem].max_allocated_mem = total_mem + total_summarized_mem; + } else { + // Pool free: end stripe, shift stripes above down within the pool. + // The envelope stays at its high-water mark (never shrinks). + const pool = pools[pool_key]; + const block_idx = pool.block_stack.findIndex(b => b.elem === elem); + if (block_idx >= 0) { + advance(1); + const block = pool.block_stack[block_idx]; + block.stripe_data.timesteps.push(timestep); + block.stripe_data.offsets.push(block.stripe_data.offsets.at(-1)); + + pool.block_stack.splice(block_idx, 1); + pool.active -= size; + + if (block_idx < pool.block_stack.length) { + for (let j = block_idx; j < pool.block_stack.length; j++) { + const b = pool.block_stack[j]; + b.inner_offset -= size; + const s = b.stripe_data; + s.timesteps.push(timestep); + s.offsets.push(s.offsets.at(-1)); + s.timesteps.push(timestep + 3); + s.offsets.push(pool.envelope_data.offsets.at(-1) + b.inner_offset); + } + advance(3); + } + } else { + pool.active -= size; + advance(1); + } + delete pool_active_elems[elem]; + } + max_size = Math.max(total_mem + total_summarized_mem, max_size); + continue; + } + + // --- Non-pool element --- + if (!(elem in draw_elem)) { + // Too small to render individually — goes into the summarized band + if (elem in summarized_elems) { + advance(1); + total_summarized_mem -= size; + summarized_elems[elem] = null; + } else { + total_summarized_mem += size; + summarized_elems[elem] = true; + advance(1); + } + continue; + } + const idx = current.findLastIndex(x => x === elem); + if (idx === -1) { + // First appearance → alloc + add_allocation(elem); + advance(1); + } else { + // Second appearance → free: remove from stack, shift elements above down + advance(1); + const removed = current_data[idx]; + removed.timesteps.push(timestep); + removed.offsets.push(removed.offsets.at(-1)); + current.splice(idx, 1); + current_data.splice(idx, 1); + + if (idx < current.length) { + shift_elements_above(idx, -size); + advance(3); + } + total_mem -= size; + } + max_size = Math.max(total_mem + total_summarized_mem, max_size); + } + + // --- Finalize: close all still-active elements --- + for (const elem of current_data) { + elem.timesteps.push(timestep); + elem.offsets.push(elem.offsets.at(-1)); + if (Array.isArray(elem.size)) { + elem.size.push(elem.size.at(-1)); + } + } + for (const pk in pools) { + for (const block of pools[pk].block_stack) { + const s = block.stripe_data; + s.timesteps.push(timestep); + s.offsets.push(s.offsets.at(-1)); + } + } + data.push(summarized_mem); + + return { + max_size, + allocations_over_time: data, + max_at_time, + summarized_mem, + elements_length: elements.length, + context_for_id: id => { + const elem = elements[id]; + let text = `Addr: ${formatAddr(elem)}`; + text = `${text}, Size: ${formatSize(elem.size)} allocation`; + text = `${text}, Total memory used after allocation: ${formatSize( + elem.max_allocated_mem, + )}`; + const context = elem?.compile_context ?? 'None'; + text = `${text}, Compile context: ${context}`; + if (elem.stream !== null) { + text = `${text}, stream ${elem.stream}`; + } + if (elem.segment_pool_id) { + text = `${text}, pool_id (${elem.segment_pool_id[0]}, ${elem.segment_pool_id[1]})`; + } else { + text = `${text}, pool_id unknown`; + } + if (elem.timestamp !== null) { + var d = new Date(elem.time_us / 1000); + text = `${text}, timestamp ${d}`; + } + if (!elem.action.includes('alloc')) { + text = `${text}\nalloc not recorded, stack trace for free:`; + } + if (elem.ghost) { + text = `${text}\n[Ghost block] This block exists in the segment snapshot but has no alloc trace events. ` + + `It was allocated before _record_memory_history() was called, or its alloc event was evicted ` + + `from the trace ring buffer. The block is still active (not freed) at snapshot time.`; + } + const user_metadata_str = format_user_metadata(elem.user_metadata); + if (user_metadata_str) { + text = `${text}\n${user_metadata_str}`; + } + text = `${text}\n${format_frames(elem.frames)}`; + text = `${text}${format_forward_frames(elem.forward_frames)}`; + return text; + }, + }; +} + +export { process_alloc_data, isPrivatePoolId, formatSize, formatAddr, + elideRepeats, frameFilter, format_user_metadata, + format_forward_frames, format_frames }; diff --git a/torchgen/aoti/fallback_ops.py b/torchgen/aoti/fallback_ops.py index 04af02ce848db..939c9fdc1cdd9 100644 --- a/torchgen/aoti/fallback_ops.py +++ b/torchgen/aoti/fallback_ops.py @@ -152,6 +152,12 @@ "aten.randn.default": {}, "aten.randn.generator": {}, "aten.randperm.default": {}, + "aten.rand_like.default": {}, + "aten.rand_like.generator": {}, + "aten.randint_like.default": {}, + "aten.randint_like.low_dtype": {}, + "aten.randn_like.default": {}, + "aten.randn_like.generator": {}, "aten.repeat_interleave.Tensor": {}, "aten.replication_pad1d_backward.default": {}, "aten.replication_pad2d_backward.default": {}, diff --git a/torchgen/api/python.py b/torchgen/api/python.py index ca971e854b234..254d7c1ee9b43 100644 --- a/torchgen/api/python.py +++ b/torchgen/api/python.py @@ -288,7 +288,7 @@ def argument_str_pyi( name += "_" # pyi merges the _out and functional variants into the same signature, with an optional out arg - if name == "out" and type_str == "Tensor" and not deprecated: + if name == "out" and not deprecated: type_str = f"{type_str} | None".replace(" | None | None", " | None") # pyi deprecated signatures don't get defaults for their out arg @@ -975,14 +975,17 @@ def argument_type_str_pyi(t: Type) -> str: if str(t.elem) == "int": ret = "_int | _size" if t.size is not None else "_size" elif t.is_tensor_like(): - # TODO: this doesn't seem right... - # Tensor?[] currently translates to tuple[Tensor, ...] | list[Tensor] | None - # It should probably translate to tuple[Tensor | None, ...] | list[Tensor | None] - add_optional = True + # Tensor?[] translates to tuple[Tensor | None, ...] | list[Tensor | None] | None + # Tensor[] translates to tuple[Tensor, ...] | list[Tensor] + if isinstance(t.elem, OptionalType): + add_optional = True + elem_str = "Tensor | None" + else: + elem_str = "Tensor" ret = ( - "Tensor | tuple[Tensor, ...] | list[Tensor]" + f"Tensor | tuple[{elem_str}, ...] | list[{elem_str}]" if t.size is not None - else "tuple[Tensor, ...] | list[Tensor]" + else f"tuple[{elem_str}, ...] | list[{elem_str}]" ) elif str(t.elem) == "float": ret = "Sequence[_float]" diff --git a/torchgen/static_runtime/generator.py b/torchgen/static_runtime/generator.py index e15b2514830ac..d8aba4d13bcde 100644 --- a/torchgen/static_runtime/generator.py +++ b/torchgen/static_runtime/generator.py @@ -270,7 +270,7 @@ def is_supported(g: NativeFunctionsGroup | NativeFunctionsViewGroup) -> bool: # the string, just test the dang thing directly if "at::Tensor" != cpp.returns_type(func.returns, symint=False).cpp_type(): # Returns a non-Tensor value. - logger.info("NON-TENSOR RET TYPE: %s", str(func)) + logger.info("NON-TENSOR RET TYPE: %s", func) return False return True