Enable AArch64 CI scripts to be used for local dev

aarch64_linux/aarch64_ci_build.sh

@@ -4,7 +4,12 @@ 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
+if [ -d /opt/conda ]; then
+  source /opt/conda/etc/profile.d/conda.sh
+  conda activate aarch64_env
+else
+  source $SCRIPTPATH/aarch64_ci_setup.sh
+fi
 
 tagged_version() {
   GIT_DESCRIBE="git --git-dir /pytorch/.git describe --tags --match v[0-9]*.[0-9]*.[0-9]*"

aarch64_linux/aarch64_wheel_ci_build.py

@@ -32,32 +32,34 @@ def build_ArmComputeLibrary() -> None:
         "build=native",
     ]
     acl_install_dir = "/acl"
-    acl_checkout_dir = "ComputeLibrary"
-    os.makedirs(acl_install_dir)
-    check_call(
-        [
-            "git",
-            "clone",
-            "https://github.com/ARM-software/ComputeLibrary.git",
-            "-b",
-            "v24.04",
-            "--depth",
-            "1",
-            "--shallow-submodules",
-        ]
-    )
+    acl_checkout_dir = "/ComputeLibrary"
+    if os.path.isdir(acl_install_dir):
+        shutil.rmtree(acl_install_dir)
+    if not os.path.isdir(acl_checkout_dir) or not len(os.listdir(acl_checkout_dir)):
+        check_call(
+            [
+                "git",
+                "clone",
+                "https://github.com/ARM-software/ComputeLibrary.git",
+                "-b",
+                "v24.08",
+                "--depth",
+                "1",
+                "--shallow-submodules",
+            ]
+        )
+        # patch Winograd conv initialzation to avoid SIGILL crash on Cortex A72
+        print("Applying ACL patch to fix SIGILL crash")
+        with open(os.path.join(os.path.dirname(__file__), "0001-Delay-Winograd-transform-initialization.patch")) as f:
+            check_call(["patch", "-p1"], stdin=f, cwd=acl_checkout_dir)
 
-    # patch Winograd conv initialzation to avoid SIGILL crash on Cortex A72
-    print("Applying ACL patch to fix SIGILL crash")
-    with open(os.path.join(os.path.dirname(__file__), "0001-Delay-Winograd-transform-initialization.patch")) as f:
-        check_call(["patch", "-p1"], stdin=f, cwd=acl_checkout_dir)
 
     check_call(
-        ["scons", "Werror=1", "-j8", f"build_dir=/{acl_install_dir}/build"]
+        ["scons", "Werror=1", f"-j{os.cpu_count()}"]
         + acl_build_flags,
         cwd=acl_checkout_dir,
     )
-    for d in ["arm_compute", "include", "utils", "support", "src"]:
+    for d in ["arm_compute", "include", "utils", "support", "src", "build"]:
         shutil.copytree(f"{acl_checkout_dir}/{d}", f"{acl_install_dir}/{d}")
 
 
@@ -181,7 +183,6 @@ def parse_arguments():
 
     print("Building PyTorch wheel")
     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: