try setting MAX_JOBS=4 for oom in arm wheel

.github/scripts/validate_binaries.sh

@@ -54,23 +54,23 @@ else
         ${PWD}/check_binary.sh
     fi
 
+     # We are only interested in CUDA tests and Python 3.8-3.11. Not all requirement libraries are available for 3.12 yet.
+    if [[ ${INCLUDE_TEST_OPS:-} == 'true' &&  ${MATRIX_GPU_ARCH_TYPE} == 'cuda' && ${MATRIX_PYTHON_VERSION} != "3.12" ]]; then
+        source ./.github/scripts/validate_test_ops.sh
+    fi
+
     if [[ ${TARGET_OS} == 'windows' ]]; then
         python  ./test/smoke_test/smoke_test.py ${TEST_SUFFIX}
     else
-        python3  ./test/smoke_test/smoke_test.py ${TEST_SUFFIX}
+        python3  ./test/smoke_test/smoke_test.py ${TEST_SUFFIX} --runtime-error-check disabled
     fi
 
     if [[ ${TARGET_OS} == 'macos-arm64' ]]; then
         export PATH=${OLD_PATH}
     fi
 
-    # We are only interested in CUDA tests and Python 3.8-3.11. Not all requirement libraries are available for 3.12 yet.
-    if [[ ${INCLUDE_TEST_OPS:-} == 'true' &&  ${MATRIX_GPU_ARCH_TYPE} == 'cuda' && ${MATRIX_PYTHON_VERSION} != "3.12" ]]; then
-        source ./.github/scripts/validate_test_ops.sh
-    fi
-
-    # TODO: remove if statement currently this step is timing out on linx-aarch64
-    if [[ ${TARGET_OS} != 'linux-aarch64' ]]; then
+    # this is optional step
+    if [[ ${TARGET_OS} != linux*  ]]; then
         conda deactivate
         conda env remove -n ${ENV_NAME}
     fi

.github/workflows/build-manywheel-images.yml

@@ -12,7 +12,9 @@ on:
     paths:
       - .github/workflows/build-manywheel-images.yml
       - manywheel/Dockerfile
+      - manywheel/Dockerfile_2_28
       - manywheel/Dockerfile_aarch64
+      - manywheel/Dockerfile_2_28_aarch64
       - manywheel/Dockerfile_cuda_aarch64
       - manywheel/Dockerfile_cxx11-abi
       - manywheel/build_docker.sh
@@ -21,7 +23,9 @@ on:
     paths:
       - .github/workflows/build-manywheel-images.yml
       - manywheel/Dockerfile
+      - manywheel/Dockerfile_2_28
       - manywheel/Dockerfile_aarch64
+      - manywheel/Dockerfile_2_28_aarch64
       - manywheel/Dockerfile_cuda_aarch64
       - manywheel/Dockerfile_cxx11-abi
       - 'common/*'
@@ -56,6 +60,27 @@ jobs:
       - name: Build Docker Image
         run: |
           manywheel/build_docker.sh
+  build-docker-cuda-manylinux_2_28:
+    runs-on: linux.12xlarge.ephemeral
+    strategy:
+      matrix:
+        cuda_version: ["12.4", "12.1", "11.8"]
+    env:
+      GPU_ARCH_TYPE: cuda-manylinux_2_28
+      GPU_ARCH_VERSION: ${{ matrix.cuda_version }}
+    steps:
+      - name: Purge tools folder (free space for build)
+        run: rm -rf /opt/hostedtoolcache
+      - name: Checkout PyTorch builder
+        uses: actions/checkout@v3
+      - name: Authenticate if WITH_PUSH
+        run: |
+          if [[ "${WITH_PUSH}" == true ]]; then
+            echo "${DOCKER_TOKEN}" | docker login -u "${DOCKER_ID}" --password-stdin
+          fi
+      - name: Build Docker Image
+        run: |
+          manywheel/build_docker.sh
   build-docker-cuda-aarch64:
     runs-on: linux.arm64.2xlarge
     strategy:
@@ -107,6 +132,21 @@ jobs:
       - name: Build Docker Image
         run: |
           manywheel/build_docker.sh
+  build-docker-cpu-manylinux_2_28:
+    runs-on: ubuntu-22.04
+    env:
+      GPU_ARCH_TYPE: cpu-manylinux_2_28
+    steps:
+      - name: Checkout PyTorch
+        uses: actions/checkout@v3
+      - name: Authenticate if WITH_PUSH
+        run: |
+          if [[ "${WITH_PUSH}" == true ]]; then
+            echo "${DOCKER_TOKEN}" | docker login -u "${DOCKER_ID}" --password-stdin
+          fi
+      - name: Build Docker Image
+        run: |
+          manywheel/build_docker.sh
   build-docker-cpu-aarch64:
     runs-on: linux.arm64.2xlarge
     env:
@@ -122,10 +162,10 @@ jobs:
       - name: Build Docker Image
         run: |
           manywheel/build_docker.sh
-  build-docker-cpu-cxx11-abi:
-    runs-on: ubuntu-22.04
+  build-docker-cpu-aarch64-2_28:
+    runs-on: linux.arm64.2xlarge
     env:
-      GPU_ARCH_TYPE: cpu-cxx11-abi
+      GPU_ARCH_TYPE: cpu-aarch64-2_28
     steps:
       - name: Checkout PyTorch
         uses: actions/checkout@v3
@@ -137,10 +177,10 @@ jobs:
       - name: Build Docker Image
         run: |
           manywheel/build_docker.sh
-  build-docker-cpu-s390x:
-    runs-on: linux.s390x
+  build-docker-cpu-cxx11-abi:
+    runs-on: ubuntu-22.04
     env:
-      GPU_ARCH_TYPE: cpu-s390x
+      GPU_ARCH_TYPE: cpu-cxx11-abi
     steps:
       - name: Checkout PyTorch
         uses: actions/checkout@v3

.github/workflows/validate-windows-binaries.yml

@@ -127,7 +127,7 @@ jobs:
 
         printf '%s\n' ${{ toJson(inputs.release-matrix) }} > release_matrix.json
         source /c/Jenkins/Miniconda3/etc/profile.d/conda.sh
-        if [[ ${MATRIX_GPU_ARCH_VERSION} == "12.1" ]]; then
+        if [[ ${MATRIX_GPU_ARCH_TYPE} == "cuda" ]]; then
           ./windows/internal/driver_update.bat
         fi
         source ./.github/scripts/validate_binaries.sh

aarch64_linux/README.md

@@ -16,4 +16,4 @@ __NOTE:__ CI build is currently __EXPERMINTAL__
 This app allows a person to build using AWS EC3 resources and requires AWS-CLI and Boto3 with AWS credentials to support building EC2 instances for the wheel builds. Can be used in a codebuild CD or from a local system.
 
 ### Usage
-```build_aarch64_wheel.py --key-name <YourPemKey> --use-docker --python 3.8 --branch <RCtag>```
+```build_aarch64_wheel.py --key-name <YourPemKey> --use-docker --python 3.8 --branch <RCtag>```

aarch64_linux/aarch64_ci_build.sh

@@ -1,6 +1,8 @@
 #!/bin/bash
 set -eux -o pipefail
 
+GPU_ARCH_VERSION=${GPU_ARCH_VERSION:-}
+
 SCRIPTPATH="$( cd -- "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
 source $SCRIPTPATH/aarch64_ci_setup.sh
 
@@ -26,10 +28,10 @@ cd /
 git config --global --add safe.directory /pytorch
 pip install -r /pytorch/requirements.txt
 pip install auditwheel
-if [ -n "$GPU_ARCH_VERSION" ]; then
-    echo "BASE_CUDA_VERSION is set to: $GPU_ARCH_VERSION"
-    python /builder/aarch64_linux/aarch64_wheel_ci_build.py --enable-mkldnn --enable-cuda
-else
-    echo "BASE_CUDA_VERSION is not set."
+if [ "$DESIRED_CUDA" = "cpu" ]; then
+    echo "BASE_CUDA_VERSION is not set. Building cpu wheel."
     python /builder/aarch64_linux/aarch64_wheel_ci_build.py --enable-mkldnn
+else
+    echo "BASE_CUDA_VERSION is set to: $DESIRED_CUDA"
+    python /builder/aarch64_linux/aarch64_wheel_ci_build.py --enable-mkldnn --enable-cuda
 fi

aarch64_linux/aarch64_wheel_ci_build.py

@@ -78,7 +78,7 @@ def build_ArmComputeLibrary() -> None:
             "clone",
             "https://github.com/ARM-software/ComputeLibrary.git",
             "-b",
-            "v23.08",
+            "v24.04",
             "--depth",
             "1",
             "--shallow-submodules",
@@ -122,12 +122,10 @@ def update_wheel(wheel_path) -> None:
         "/usr/local/cuda/lib64/libcudnn_cnn_train.so.8",
         "/usr/local/cuda/lib64/libcudnn_ops_infer.so.8",
         "/usr/local/cuda/lib64/libcudnn_ops_train.so.8",
-        "/opt/conda/envs/aarch64_env/lib/libopenblas.so.0",
-        "/opt/conda/envs/aarch64_env/lib/libgfortran.so.5",
         "/opt/conda/envs/aarch64_env/lib/libgomp.so.1",
+        "/opt/OpenBLAS/lib/libopenblas.so.0",
         "/acl/build/libarm_compute.so",
         "/acl/build/libarm_compute_graph.so",
-        "/acl/build/libarm_compute_core.so",
     ]
     # Copy libraries to unzipped_folder/a/lib
     for lib_path in libs_to_copy:
@@ -140,10 +138,10 @@ def update_wheel(wheel_path) -> None:
     os.system(f"cd {folder}/tmp/; zip -r {folder}/cuda_wheel/{wheelname} *")
     shutil.move(
         f"{folder}/cuda_wheel/{wheelname}",
-        f"/dist/{wheelname}",
+        f"{folder}/{wheelname}",
         copy_function=shutil.copy2,
     )
-    os.system(f"rm -rf {folder}/tmp {folder}/dist/cuda_wheel/")
+    os.system(f"rm -rf {folder}/tmp/ {folder}/cuda_wheel/")
 
 
 def complete_wheel(folder: str) -> str:
@@ -201,8 +199,9 @@ def parse_arguments():
         branch = "master"
 
     print("Building PyTorch wheel")
-    build_vars = "CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
-    os.system("python setup.py clean")
+    os.system("export USE_PRIORITIZED_TEXT_FOR_LD=1")
+    build_vars = "MAX_JOBS=5 CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
+    os.system("cd /pytorch; python setup.py clean")
 
     override_package_version = os.getenv("OVERRIDE_PACKAGE_VERSION")
     if override_package_version is not None:

aarch64_linux/build_aarch64_wheel.py

@@ -229,7 +229,7 @@ def build_ArmComputeLibrary(host: RemoteHost, git_clone_flags: str = "") -> None
     print('Building Arm Compute Library')
     acl_build_flags=" ".join(["debug=0", "neon=1", "opencl=0", "os=linux", "openmp=1", "cppthreads=0",
                               "arch=armv8a", "multi_isa=1", "fixed_format_kernels=1", "build=native"])
-    host.run_cmd(f"git clone https://github.com/ARM-software/ComputeLibrary.git -b v23.08 {git_clone_flags}")
+    host.run_cmd(f"git clone https://github.com/ARM-software/ComputeLibrary.git -b v24.04 {git_clone_flags}")
     host.run_cmd(f"cd ComputeLibrary && scons Werror=1 -j8 {acl_build_flags}")
 
 

check_binary.sh

@@ -330,7 +330,7 @@ fi
 if [[ "$PACKAGE_TYPE" == 'libtorch' ]]; then
   echo "Checking that MKL is available"
   build_and_run_example_cpp check-torch-mkl
-elif [[ "$(uname -m)" != "arm64" ]]; then
+elif [[ "$(uname -m)" != "arm64" && "$(uname -m)" != "s390x" ]]; then
   if [[ "$(uname)" != 'Darwin' || "$PACKAGE_TYPE" != *wheel ]]; then
     if [[ "$(uname -m)" == "aarch64" ]]; then
       echo "Checking that MKLDNN is available on aarch64"
@@ -354,7 +354,7 @@ if [[ "$PACKAGE_TYPE" == 'libtorch' ]]; then
   echo "Checking that XNNPACK is available"
   build_and_run_example_cpp check-torch-xnnpack
 else
-  if [[ "$(uname)" != 'Darwin' || "$PACKAGE_TYPE" != *wheel ]]; then
+  if [[ "$(uname)" != 'Darwin' || "$PACKAGE_TYPE" != *wheel ]] && [[ "$(uname -m)" != "s390x"  ]]; then
     echo "Checking that XNNPACK is available"
     pushd /tmp
     python -c 'import torch.backends.xnnpack; exit(0 if torch.backends.xnnpack.enabled else 1)'
@@ -375,7 +375,7 @@ if [[ "$OSTYPE" == "msys" ]]; then
 fi
 
 # Test that CUDA builds are setup correctly
-if [[ "$DESIRED_CUDA" != 'cpu' && "$DESIRED_CUDA" != 'cpu-cxx11-abi' && "$DESIRED_CUDA" != *"rocm"* ]]; then
+if [[ "$DESIRED_CUDA" != 'cpu' && "$DESIRED_CUDA" != 'cpu-cxx11-abi' && "$DESIRED_CUDA" != *"rocm"* && "$(uname -m)" != "s390x" ]]; then
   if [[ "$PACKAGE_TYPE" == 'libtorch' ]]; then
     build_and_run_example_cpp check-torch-cuda
   else