From 2535a25ed99ddb3bb94a96b37d8fc6c9aee209cf Mon Sep 17 00:00:00 2001
From: Hansong Zhang <hsz@meta.com>
Date: Tue, 31 Oct 2023 17:50:36 -0700
Subject: [PATCH] Arm build and runner CI (#1111)

Summary:
Add the setup part to docker image to make the run faster.

Pull Request resolved: https://github.com/pytorch/executorch/pull/1111

Reviewed By: huydhn

Differential Revision: D50866420

Pulled By: kirklandsign

fbshipit-source-id: 6c2e6ff7ffe245df950519dd258204661e2fed56
---
 .ci/docker/build.sh                 |  9 ++++++++-
 .ci/docker/ubuntu/Dockerfile        |  5 +++++
 .github/workflows/docker-builds.yml |  1 +
 .github/workflows/trunk.yml         | 18 ++++++++++++++++++
 backends/arm/arm_backend.py         |  4 +---
 5 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/.ci/docker/build.sh b/.ci/docker/build.sh
index 1c4e2a72f5..48d95f8b45 100755
--- a/.ci/docker/build.sh
+++ b/.ci/docker/build.sh
@@ -26,6 +26,9 @@ case "${IMAGE_NAME}" in
   executorch-ubuntu-22.04-linter)
     LINTRUNNER=yes
     ;;
+  executorch-ubuntu-22.04-arm-sdk)
+    ARM_SDK=yes
+    ;;
   *)
     echo "Invalid image name ${IMAGE_NAME}"
     exit 1
@@ -41,6 +44,9 @@ BUILD_DOCS=1
 # Copy requirements-lintrunner.txt from root to here
 cp ../../requirements-lintrunner.txt ./
 
+# Copy arm setup script from root to here
+cp -r ../../examples/arm/ ./arm
+
 docker build \
   --no-cache \
   --progress=plain \
@@ -52,8 +58,9 @@ docker build \
   --build-arg "TORCHAUDIO_VERSION=${TORCHAUDIO_VERSION}.${NIGHTLY}" \
   --build-arg "TORCHVISION_VERSION=${TORCHVISION_VERSION}.${NIGHTLY}" \
   --build-arg "BUCK2_VERSION=${BUCK2_VERSION}" \
-  --build-arg "LINTRUNNER=${LINTRUNNER}" \
+  --build-arg "LINTRUNNER=${LINTRUNNER:-}" \
   --build-arg "BUILD_DOCS=${BUILD_DOCS}" \
+  --build-arg "ARM_SDK=${ARM_SDK:-}" \
   -f "${OS}"/Dockerfile \
   "$@" \
   .
diff --git a/.ci/docker/ubuntu/Dockerfile b/.ci/docker/ubuntu/Dockerfile
index 6afb6ab7bd..cd9a0a1f37 100644
--- a/.ci/docker/ubuntu/Dockerfile
+++ b/.ci/docker/ubuntu/Dockerfile
@@ -51,5 +51,10 @@ COPY ./requirements-lintrunner.txt requirements-lintrunner.txt
 RUN if [ -n "${LINTRUNNER}" ]; then bash ./install_linter.sh; fi
 RUN rm install_linter.sh utils.sh requirements-lintrunner.txt
 
+ARG ARM_SDK
+COPY --chown=ci-user:ci-user ./arm /opt/arm
+# Set up ARM SDK if needed
+RUN if [ -n "${ARM_SDK}" ]; then git config --global user.email "you@example.com"; git config --global user.name "OSS CI"; bash /opt/arm/setup.sh --i-agree-to-the-contained-eula /opt/arm-sdk; chown -R ci-user:ci-user /opt/arm-sdk; fi
+
 USER ci-user
 CMD ["bash"]
diff --git a/.github/workflows/docker-builds.yml b/.github/workflows/docker-builds.yml
index 814961a5b9..07bf7150ec 100644
--- a/.github/workflows/docker-builds.yml
+++ b/.github/workflows/docker-builds.yml
@@ -27,6 +27,7 @@ jobs:
         include:
           - docker-image-name: executorch-ubuntu-22.04-clang12
           - docker-image-name: executorch-ubuntu-22.04-linter
+          - docker-image-name: executorch-ubuntu-22.04-arm-sdk
     env:
       DOCKER_IMAGE: 308535385114.dkr.ecr.us-east-1.amazonaws.com/executorch/${{ matrix.docker-image-name }}
     steps:
diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml
index 63625f76eb..047c59525a 100644
--- a/.github/workflows/trunk.yml
+++ b/.github/workflows/trunk.yml
@@ -131,6 +131,24 @@ jobs:
         # Test selective build
         PYTHON_EXECUTABLE=python bash examples/portable/scripts/test_demo_backend_delegation.sh "${BUILD_TOOL}"
 
+  test-arm-backend-delegation:
+    name: test-arm-backend-delegation
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      runner: linux.2xlarge
+      docker-image: executorch-ubuntu-22.04-arm-sdk
+      submodules: 'true'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      script: |
+        # The generic Linux job chooses to use base env, not the one setup by the image
+        CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]")
+        conda activate "${CONDA_ENV}"
+
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh cmake
+        # Test selective build
+        source /opt/arm-sdk/setup_path.sh
+        PYTHON_EXECUTABLE=python bash examples/arm/run.sh /opt/arm-sdk buck2
+
   test-coreml-delegate:
     name: test-coreml-delegate
     uses: pytorch/test-infra/.github/workflows/macos_job.yml@main
diff --git a/backends/arm/arm_backend.py b/backends/arm/arm_backend.py
index 2a4fca05e6..fc2fe45c1a 100644
--- a/backends/arm/arm_backend.py
+++ b/backends/arm/arm_backend.py
@@ -196,9 +196,7 @@ def vela_compile(tosa_fb):
             # Add a block for scratch, inputs and outputs;  scratch shape is a 1 element
             # array giving us size in bytes so extract this and add a block of 0's.
             # Currently we preallocated this on the host to provide SRAM for computation.
-            if len(data["scratch_shape"][0]) != 1:
-                raise RuntimeError("Expected scratch to be single array")
-            block_length = data["scratch_shape"][0].item()
+            block_length = int(data["scratch_shape"][0])
             bin_blocks["scratch_data"] = b"\x00" * block_length
 
             # Capture inputs and outputs