Improve performance of reduce sum for 3D shapes

.ci/docker/build.sh

@@ -50,6 +50,8 @@ if [[ "$image" == *-focal* ]]; then
   UBUNTU_VERSION=20.04
 elif [[ "$image" == *-jammy* ]]; then
   UBUNTU_VERSION=22.04
+elif [[ "$image" == *-noble* ]]; then
+  UBUNTU_VERSION=24.04
 elif [[ "$image" == *ubuntu* ]]; then
   extract_version_from_image_name ubuntu UBUNTU_VERSION
 elif [[ "$image" == *centos* ]]; then
@@ -452,10 +454,15 @@ if [[ "$image" == *cuda*  && ${OS} == "ubuntu" ]]; then
   fi
 fi
 
+DOCKER_PROGRESS="--progress=plain"
+if [[ "${DOCKER_BUILDKIT}" == 0 ]]; then
+  DOCKER_PROGRESS=""
+fi
+
 # Build image
 docker build \
        --no-cache \
-       --progress=plain \
+       ${DOCKER_PROGRESS} \
        --build-arg "BUILD_ENVIRONMENT=${image}" \
        --build-arg "PROTOBUF=${PROTOBUF:-}" \
        --build-arg "LLVMDEV=${LLVMDEV:-}" \

.ci/docker/centos-rocm/Dockerfile

@@ -80,6 +80,8 @@ RUN rm install_rocm_magma.sh
 COPY ./common/install_amdsmi.sh install_amdsmi.sh
 RUN bash ./install_amdsmi.sh
 RUN rm install_amdsmi.sh
+
+ENV ROCM_PATH /opt/rocm
 ENV PATH /opt/rocm/bin:$PATH
 ENV PATH /opt/rocm/hcc/bin:$PATH
 ENV PATH /opt/rocm/hip/bin:$PATH
@@ -113,13 +115,6 @@ COPY triton_version.txt triton_version.txt
 RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
 RUN rm install_triton.sh common_utils.sh triton-rocm.txt triton_version.txt
 
-# Install AOTriton (Early fail)
-COPY ./aotriton_version.txt aotriton_version.txt
-COPY ./common/common_utils.sh common_utils.sh
-COPY ./common/install_aotriton.sh install_aotriton.sh
-RUN ["/bin/bash", "-c", "./install_aotriton.sh /opt/rocm && rm -rf install_aotriton.sh aotriton_version.txt common_utils.sh"]
-ENV AOTRITON_INSTALLED_PREFIX /opt/rocm/aotriton
-
 # Install ccache/sccache (do this last, so we get priority in PATH)
 COPY ./common/install_cache.sh install_cache.sh
 ENV PATH /opt/cache/bin:$PATH

.ci/docker/ci_commit_pins/triton-rocm.txt

@@ -1 +1 @@
-21eae954efa5bf584da70324b640288c3ee7aede
+75cc27c26a88b4edbcd11671a8aa524b65478d46

.ci/docker/common/cache_vision_models.sh

@@ -2,6 +2,20 @@
 
 set -ex
 
+# Skip pytorch-nightly installation in docker images
+# Installation of pytorch-nightly is needed to prefetch mobilenet_v2 avd v3 models for some tests.
+# Came from https://github.com/ROCm/pytorch/commit/85bd6bc0105162293fa0bbfb7b661f85ec67f85a
+# Models are downloaded on first use to the folder /root/.cache/torch/hub
+# But pytorch-nightly installation also overrides .ci/docker/requirements-ci.txt settings
+# and upgrades some of python packages (sympy from 1.12.0 to 1.13.0)
+# which causes several 'dynamic_shapes' tests to fail
+# Skip prefetching models affects these tests without any errors:
+#   python test/mobile/model_test/gen_test_model.py mobilenet_v2
+#   python test/quantization/eager/test_numeric_suite_eager.py -k test_mobilenet_v3
+# Issue https://github.com/ROCm/frameworks-internal/issues/8772
+echo "Skip torch-nightly installation"
+exit 0
+
 source "$(dirname "${BASH_SOURCE[0]}")/common_utils.sh"
 
 # Cache the test models at ~/.cache/torch/hub/

.ci/docker/common/common_utils.sh

@@ -23,6 +23,10 @@ conda_install() {
   as_jenkins conda install -q -n py_$ANACONDA_PYTHON_VERSION -y python="$ANACONDA_PYTHON_VERSION" $*
 }
 
+conda_install_through_forge() {
+  as_jenkins conda install -c conda-forge -q -n py_$ANACONDA_PYTHON_VERSION -y python="$ANACONDA_PYTHON_VERSION" $*
+}
+
 conda_run() {
   as_jenkins conda run -n py_$ANACONDA_PYTHON_VERSION --no-capture-output $*
 }

.ci/docker/common/install_base.sh

@@ -15,6 +15,9 @@ install_ubuntu() {
   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"
@@ -82,11 +85,42 @@ install_ubuntu() {
   # see: https://github.com/pytorch/pytorch/issues/65931
   apt-get install -y libgnutls30
 
+  # Required to install the fortran after gcc update
+  if [[ "$UBUNTU_VERSION" == "22.04"* ]]; then
+    apt autoremove -y gfortran
+    apt-get update -y
+    apt-get install -y gfortran libopenblas-dev
+  fi
+
   # Cleanup package manager
   apt-get autoclean && apt-get clean
   rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 }
 
+build_libpng() {
+  # install few packages
+  yum install -y zlib zlib-devel
+
+  LIBPNG_VERSION=1.6.37
+
+  mkdir -p libpng
+  pushd libpng
+
+  wget http://download.sourceforge.net/libpng/libpng-$LIBPNG_VERSION.tar.gz
+  tar -xvzf libpng-$LIBPNG_VERSION.tar.gz
+
+  pushd libpng-$LIBPNG_VERSION
+
+  ./configure
+  make
+  make install
+
+  popd
+
+  popd
+  rm -rf libpng
+}
+
 install_centos() {
   # Need EPEL for many packages we depend on.
   # See http://fedoraproject.org/wiki/EPEL
@@ -123,6 +157,11 @@ install_centos() {
     unzip \
     gdb
 
+  # CentOS7 doesnt have support for higher version of libpng,
+  # so it is built from source.
+  # Libpng is required for torchvision build.
+  build_libpng
+
   # Cleanup
   yum clean all
   rm -rf /var/cache/yum

.ci/docker/common/install_conda.sh

@@ -38,7 +38,10 @@ fi
   source "$(dirname "${BASH_SOURCE[0]}")/common_utils.sh"
 
   pushd /tmp
-  wget -q "${BASE_URL}/${CONDA_FILE}"
+  if [ -n $CENTOS_VERSION ] && [[ $CENTOS_VERSION == 7.* ]]; then
+    NO_CHECK_CERTIFICATE_FLAG="--no-check-certificate"
+  fi
+  wget -q "${BASE_URL}/${CONDA_FILE}" ${NO_CHECK_CERTIFICATE_FLAG}
   # NB: Manually invoke bash per https://github.com/conda/conda/issues/10431
   as_jenkins bash "${CONDA_FILE}" -b -f -p "/opt/conda"
   popd
@@ -110,6 +113,15 @@ fi
     conda_install magma-cuda$(TMP=${CUDA_VERSION/./};echo ${TMP%.*[0-9]}) -c pytorch
   fi
 
+  # Install required libstdc++.so.6 version
+  if [ "$ANACONDA_PYTHON_VERSION" = "3.10" ] || [ "$ANACONDA_PYTHON_VERSION" = "3.9" ] ; then
+    conda_install_through_forge libstdcxx-ng=12
+  fi
+
+  if [ "$ANACONDA_PYTHON_VERSION" = "3.12" ] || [ "$UBUNTU_VERSION" == "24.04"* ] ; then
+    conda_install_through_forge libstdcxx-ng=14
+  fi
+
   # Install some other packages, including those needed for Python test reporting
   pip_install -r /opt/conda/requirements-ci.txt
 

.ci/docker/common/install_rocm.sh

@@ -16,6 +16,11 @@ install_ubuntu() {
       # gpg-agent is not available by default on 20.04
       apt-get install -y --no-install-recommends gpg-agent
     fi
+    if [[ $UBUNTU_VERSION == 22.04 ]] || [[ $UBUNTU_VERSION == 24.04 ]]; then
+        apt-get install -y --no-install-recommends gpg-agent
+        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' \
+            | sudo tee /etc/apt/preferences.d/rocm-pin-600
+    fi
     apt-get install -y kmod
     apt-get install -y wget
 

.ci/docker/common/install_user.sh

@@ -2,6 +2,13 @@
 
 set -ex
 
+# Since version 24 the system ships with user 'ubuntu' that has id 1000
+# We need a work-around to enable id 1000 usage for this script
+if [[ $UBUNTU_VERSION == 24.04 ]]; then
+    # touch is used to disable harmless error message
+    touch /var/mail/ubuntu && chown ubuntu /var/mail/ubuntu && userdel -r ubuntu
+fi
+
 # Mirror jenkins user in container
 # jenkins user as ec2-user should have the same user-id
 echo "jenkins:x:1000:1000::/var/lib/jenkins:" >> /etc/passwd

.ci/docker/requirements-ci.txt

@@ -15,7 +15,14 @@ click
 #Pinned versions:
 #test that import:
 
+sympy==1.12.1
+#Description: Python library for symbolic mathematics
+# installed before coremltools to avoid installation of greater sympy version
+#Pinned versions: 1.12.1
+#test that import:
+
 coremltools==5.0b5 ; python_version < "3.12"
+coremltools==7.2 ; python_version == "3.12"
 #Description: Apple framework for ML integration
 #Pinned versions: 5.0b5
 #test that import:
@@ -58,6 +65,7 @@ lark==0.12.0
 #test that import:
 
 librosa>=0.6.2 ; python_version < "3.11"
+librosa==0.10.2 ; python_version == "3.12"
 #Description: A python package for music and audio analysis
 #Pinned versions: >=0.6.2
 #test that import: test_spectral_ops.py
@@ -106,6 +114,7 @@ networkx==2.8.8
 numba==0.49.0 ; python_version < "3.9"
 numba==0.54.1 ; python_version == "3.9"
 numba==0.55.2 ; python_version == "3.10"
+numba==0.60.0 ; python_version == "3.12"
 #Description: Just-In-Time Compiler for Numerical Functions
 #Pinned versions: 0.54.1, 0.49.0, <=0.49.1
 #test that import: test_numba_integration.py
@@ -247,6 +256,11 @@ tb-nightly==2.13.0a20230426
 #Pinned versions:
 #test that import:
 
+tlparse==0.3.7
+#Description: parse logs produced by torch.compile
+#Pinned versions:
+#test that import: dynamo/test_structured_trace.py
+
 # needed by torchgen utils
 typing-extensions
 #Description: type hints for python
@@ -306,7 +320,7 @@ pywavelets==1.5.0 ; python_version >= "3.12"
 #Pinned versions: 1.4.1
 #test that import:
 
-lxml==5.0.0.
+lxml==5.0.0
 #Description: This is a requirement of unittest-xml-reporting
 
 # Python-3.9 binaries

.ci/docker/ubuntu-rocm/Dockerfile

@@ -105,13 +105,6 @@ COPY triton_version.txt triton_version.txt
 RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
 RUN rm install_triton.sh common_utils.sh triton-rocm.txt triton_version.txt
 
-# Install AOTriton
-COPY ./aotriton_version.txt aotriton_version.txt
-COPY ./common/common_utils.sh common_utils.sh
-COPY ./common/install_aotriton.sh install_aotriton.sh
-RUN ["/bin/bash", "-c", "./install_aotriton.sh /opt/rocm && rm -rf install_aotriton.sh aotriton_version.txt common_utils.sh"]
-ENV AOTRITON_INSTALLED_PREFIX /opt/rocm/aotriton
-
 # 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

.circleci/scripts/binary_populate_env.sh

@@ -5,7 +5,9 @@ export TZ=UTC
 tagged_version() {
   GIT_DIR="${workdir}/pytorch/.git"
   GIT_DESCRIBE="git --git-dir ${GIT_DIR} describe --tags --match v[0-9]*.[0-9]*.[0-9]*"
-  if [[ ! -d "${GIT_DIR}" ]]; then
+  if [[ -n "${CIRCLE_TAG:-}" ]]; then
+    echo "${CIRCLE_TAG}"
+  elif [[ ! -d "${GIT_DIR}" ]]; then
     echo "Abort, abort! Git dir ${GIT_DIR} does not exists!"
     kill $$
   elif ${GIT_DESCRIBE} --exact >/dev/null; then
@@ -71,34 +73,6 @@ fi
 
 export PYTORCH_BUILD_NUMBER=1
 
-# Set triton version as part of PYTORCH_EXTRA_INSTALL_REQUIREMENTS
-TRITON_VERSION=$(cat $PYTORCH_ROOT/.ci/docker/triton_version.txt)
-
-# Here PYTORCH_EXTRA_INSTALL_REQUIREMENTS is already set for the all the wheel builds hence append TRITON_CONSTRAINT
-TRITON_CONSTRAINT="platform_system == 'Linux' and platform_machine == 'x86_64' and python_version < '3.13'"
-if [[ "$PACKAGE_TYPE" =~ .*wheel.* &&  -n "${PYTORCH_EXTRA_INSTALL_REQUIREMENTS:-}" ]]; then
-  TRITON_REQUIREMENT="triton==${TRITON_VERSION}; ${TRITON_CONSTRAINT}"
-  if [[ -n "$PYTORCH_BUILD_VERSION" && "$PYTORCH_BUILD_VERSION" =~ .*dev.* ]]; then
-      TRITON_SHORTHASH=$(cut -c1-10 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton.txt)
-      TRITON_REQUIREMENT="pytorch-triton==${TRITON_VERSION}+${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
-  fi
-  export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | ${TRITON_REQUIREMENT}"
-fi
-
-# Set triton via PYTORCH_EXTRA_INSTALL_REQUIREMENTS for triton rocm package
-if [[ "$PACKAGE_TYPE" =~ .*wheel.* && -n "$PYTORCH_BUILD_VERSION" && "$PYTORCH_BUILD_VERSION" =~ .*rocm.* && $(uname) == "Linux" ]]; then
-    TRITON_REQUIREMENT="pytorch-triton-rocm==${TRITON_VERSION}; ${TRITON_CONSTRAINT}"
-    if [[ -n "$PYTORCH_BUILD_VERSION" && "$PYTORCH_BUILD_VERSION" =~ .*dev.* ]]; then
-        TRITON_SHORTHASH=$(cut -c1-10 $PYTORCH_ROOT/.ci/docker/ci_commit_pins/triton-rocm.txt)
-        TRITON_REQUIREMENT="pytorch-triton-rocm==${TRITON_VERSION}+${TRITON_SHORTHASH}; ${TRITON_CONSTRAINT}"
-    fi
-    if [[ -z "${PYTORCH_EXTRA_INSTALL_REQUIREMENTS:-}" ]]; then
-        export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${TRITON_REQUIREMENT}"
-    else
-        export PYTORCH_EXTRA_INSTALL_REQUIREMENTS="${PYTORCH_EXTRA_INSTALL_REQUIREMENTS} | ${TRITON_REQUIREMENT}"
-    fi
-fi
-
 JAVA_HOME=
 BUILD_JNI=OFF
 if [[ "$PACKAGE_TYPE" == libtorch ]]; then

.github/scripts/amd/package_triton_wheel.sh

@@ -30,7 +30,12 @@ fi
 # Remove packaged libs and headers
 rm -rf $TRITON_ROCM_DIR/include/*
 
-LIBTINFO_PATH="/usr/lib64/libtinfo.so.5"
+OS_NAME=`awk -F= '/^NAME/{print $2}' /etc/os-release`
+if [[ "$OS_NAME" == *"CentOS Linux"* ]]; then
+    LIBTINFO_PATH="/usr/lib64/libtinfo.so.5"
+else
+    LIBTINFO_PATH="/usr/lib64/libtinfo.so.6"
+fi
 LIBNUMA_PATH="/usr/lib64/libnuma.so.1"
 LIBELF_PATH="/usr/lib64/libelf.so.1"
 
@@ -45,16 +50,9 @@ do
     cp $lib $TRITON_ROCM_DIR/lib/
 done
 
-# Required ROCm libraries
-if [[ "${MAJOR_VERSION}" == "6" ]]; then
-    libamdhip="libamdhip64.so.6"
-else
-    libamdhip="libamdhip64.so.5"
-fi
-
 # Required ROCm libraries - ROCm 6.0
 ROCM_SO=(
-    "${libamdhip}"
+    "libamdhip64.so.6"
     "libhsa-runtime64.so.1"
     "libamd_comgr.so.2"
     "libdrm.so.2"