Update on "[Executorch][llama] Make RoPE freq calculation broadcast for per head"

This is a workaround, may not be even worth landing, to avoid broadcasting
semantics in the mul op and for that matter any binary op. Current
implementation of oiptimized ops doesnt handle broadcasting and falls back to
portable op implementation.

This diff also fixes an issue where (as seen in llama) two tensors of binary op
are not broadcasting, but they have different # of dims, which results in
invocation of unoptimized path. e.g. a = [1, 1, 2048], b = [2048], out = [1, 1,
2048].

In llama case this is optimized path when generating one token at a time. Not
so during pre-fill

Making optimized op handle broadcasting, and support vectorization, is not
hard, but may take some time.

Differential Revision: [D54766067](https://our.internmc.facebook.com/intern/diff/D54766067/)

[ghstack-poisoned]

Author: kimishpatel

.ci/docker/build.sh

@@ -37,6 +37,12 @@ case "${IMAGE_NAME}" in
     ARM_SDK=yes
     CLANG_VERSION=12
     ;;
+  executorch-ubuntu-22.04-clang12-android)
+    LINTRUNNER=""
+    CLANG_VERSION=12
+    # From https://developer.android.com/ndk/downloads
+    ANDROID_NDK_VERSION=r26c
+    ;;
   *)
     echo "Invalid image name ${IMAGE_NAME}"
     exit 1
@@ -66,6 +72,7 @@ docker build \
   --build-arg "LINTRUNNER=${LINTRUNNER:-}" \
   --build-arg "BUILD_DOCS=${BUILD_DOCS}" \
   --build-arg "ARM_SDK=${ARM_SDK:-}" \
+  --build-arg "ANDROID_NDK_VERSION=${ANDROID_NDK_VERSION:-}" \
   -f "${OS}"/Dockerfile \
   "$@" \
   .

.ci/docker/common/install_android.sh

@@ -0,0 +1,70 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+set -ex
+
+# Double check if the NDK version is set
+[ -n "${ANDROID_NDK_VERSION}" ]
+
+install_prerequiresites() {
+  apt-get update
+
+  # NB: Need OpenJDK 17 at the minimum
+  apt-get install -y --no-install-recommends \
+    openjdk-17-jdk \
+    ca-certificates-java \
+    ant
+
+  # Cleanup package manager
+  apt-get autoclean && apt-get clean
+  rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+}
+
+install_ndk() {
+  NDK_INSTALLATION_DIR=/opt/ndk
+  mkdir -p "${NDK_INSTALLATION_DIR}"
+
+  pushd /tmp
+  # The NDK installation is cached on ossci-android S3 bucket
+  curl -Os --retry 3 "https://ossci-android.s3.amazonaws.com/android-ndk-${ANDROID_NDK_VERSION}-linux.zip"
+  unzip -qo "android-ndk-${ANDROID_NDK_VERSION}-linux.zip"
+
+  # Print the content for manual verification
+  ls -lah "android-ndk-${ANDROID_NDK_VERSION}"
+  mv "android-ndk-${ANDROID_NDK_VERSION}"/* "${NDK_INSTALLATION_DIR}"
+
+  popd
+}
+
+install_cmdtools() {
+  CMDTOOLS_FILENAME=commandlinetools-linux-11076708_latest.zip
+
+  pushd /tmp
+  # The file is cached on ossci-android S3 bucket
+  curl -Os --retry 3 "https://ossci-android.s3.us-west-1.amazonaws.com/${CMDTOOLS_FILENAME}"
+  unzip -qo "${CMDTOOLS_FILENAME}" -d /opt
+
+  ls -lah /opt/cmdline-tools/bin
+  popd
+}
+
+install_sdk() {
+  SDK_INSTALLATION_DIR=/opt/android/sdk
+  mkdir -p "${SDK_INSTALLATION_DIR}"
+
+  # These are the tools needed to build Android apps
+  yes | /opt/cmdline-tools/bin/sdkmanager --sdk_root="${SDK_INSTALLATION_DIR}" --install "platforms;android-34"
+  yes | /opt/cmdline-tools/bin/sdkmanager --sdk_root="${SDK_INSTALLATION_DIR}" --install "build-tools;33.0.1"
+  # And some more tools for future emulator tests
+  yes | /opt/cmdline-tools/bin/sdkmanager --sdk_root="${SDK_INSTALLATION_DIR}" --install "platform-tools"
+  yes | /opt/cmdline-tools/bin/sdkmanager --sdk_root="${SDK_INSTALLATION_DIR}" --install "tools"
+}
+
+install_prerequiresites
+install_ndk
+install_cmdtools
+install_sdk

.ci/docker/ubuntu/Dockerfile

@@ -71,6 +71,12 @@ 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 ANDROID_NDK_VERSION
+# Install lintrunner if needed
+COPY ./common/install_android.sh install_android.sh
+RUN if [ -n "${ANDROID_NDK_VERSION}" ]; then bash ./install_android.sh; fi
+RUN rm install_android.sh
+
 ARG ARM_SDK
 COPY --chown=ci-user:ci-user ./arm /opt/arm
 # Set up ARM SDK if needed

.github/workflows/android.yml

@@ -0,0 +1,47 @@
+name: Build ExecuTorch Android demo apps
+
+on:
+  push:
+    branches:
+      - main
+      - release/*
+  pull_request:
+    paths:
+      - .ci/docker/**
+      - .github/workflows/android.yml
+      - install_requirements.sh
+      - examples/demo-apps/**
+      - extension/module/**
+  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:
+  test-demo-android:
+    name: test-demo-android
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    strategy:
+      matrix:
+        include:
+          - build-tool: buck2
+    with:
+      # NB: The example model dl3 requires lots of memory (T161064121)
+      runner: linux.12xlarge
+      docker-image: executorch-ubuntu-22.04-clang12-android
+      submodules: 'true'
+      ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}
+      timeout: 90
+      script: |
+        set -eux
+
+        # 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}"
+
+        BUILD_TOOL=${{ matrix.build-tool }}
+        # Setup MacOS dependencies as there is no Docker support on MacOS atm
+        PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
+        # Build Android demo app
+        bash build/test_android_ci.sh

.github/workflows/apple.yml

@@ -5,6 +5,17 @@ on:
     branches:
       - main
       - release/*
+    paths:
+      - .ci/docker/**
+      - .github/workflows/apple.yml
+      - install_requirements.sh
+      - backends/apple/**
+      - build/build_apple_frameworks.sh
+      - build/create_frameworks.sh
+      - build/test_ios_ci.sh
+      - examples/demo-apps/**
+      - extension/apple/**
+      - extension/module/**
   pull_request:
     paths:
       - .ci/docker/**
@@ -92,9 +103,9 @@ jobs:
         build/build_apple_frameworks.sh --coreml --mps --portable --xnnpack
 
         # Bundle iOS Frameworks
-        for FRAMEWORK in "${FRAMEWORKS[@]}"; do
-          zip -r "${RUNNER_TEMP}/artifacts/${FRAMEWORK}-${VERSION}.zip" "cmake-out/${FRAMEWORK}.xcframework" LICENSE
-        done
+        for FRAMEWORK in "${FRAMEWORKS[@]}"; do (
+          cd cmake-out && zip -r "${RUNNER_TEMP}/artifacts/${FRAMEWORK}-${VERSION}.zip" "${FRAMEWORK}.xcframework"
+        ) done
 
         popd
 
@@ -122,8 +133,8 @@ jobs:
           # NB: The name here needs to match the upload-artifact name from build-frameworks-ios job
           name: executorch-frameworks-ios
           path: ${{ runner.temp }}/frameworks-ios/
-      - name: Only push to S3 from main branch
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
+      - name: Only push to S3 when running the workflow manually from main branch
+        if: ${{ github.event_name == 'workflow_dispatch' && github.ref == 'refs/heads/main' }}
         shell: bash
         run: |
           set -eux

.github/workflows/docker-builds.yml

@@ -38,6 +38,7 @@ jobs:
           - 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
+          - docker-image-name: executorch-ubuntu-22.04-clang12-android
     env:
       DOCKER_IMAGE: 308535385114.dkr.ecr.us-east-1.amazonaws.com/executorch/${{ matrix.docker-image-name }}
     steps:

.gitignore

@@ -3,6 +3,7 @@ buck-out/
 cmake-out/
 cmake-android-out/
 cmake-ios-out/
+ethos-u-scratch/
 executorch.egg-info
 __pycache__/
 build/lib/

CMakeLists.txt

@@ -152,6 +152,10 @@ option(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL
 option(EXECUTORCH_BUILD_GTESTS
        "Build googletest based test binaries" OFF)
 
+# Build the AOT util extension library
+option(EXECUTORCH_BUILD_EXTENSION_AOT_UTIL
+       "Build the extension/aot_util directory" OFF)
+
 if(NOT BUCK2)
   set(BUCK2 buck2)
 endif()
@@ -236,11 +240,11 @@ if(EXECUTORCH_BUILD_FLATC)
     )
   endif()
   set(FLATC_EXECUTABLE flatc)
-  option(FLATBUFFERS_BUILD_FLATC "" ON)
-  option(FLATBUFFERS_BUILD_FLATHASH "" OFF)
-  option(FLATBUFFERS_BUILD_FLATLIB "" OFF)
-  option(FLATBUFFERS_BUILD_TESTS "" OFF)
-  option(FLATBUFFERS_INSTALL "" OFF)
+  set(FLATBUFFERS_BUILD_FLATC ON CACHE BOOL "")
+  set(FLATBUFFERS_BUILD_FLATHASH OFF CACHE BOOL "")
+  set(FLATBUFFERS_BUILD_FLATLIB OFF CACHE BOOL "")
+  set(FLATBUFFERS_BUILD_TESTS OFF CACHE BOOL "")
+  set(FLATBUFFERS_INSTALL OFF CACHE BOOL "")
   add_subdirectory(third-party/flatbuffers)
 endif()
 if(NOT FLATC_EXECUTABLE)
@@ -358,6 +362,10 @@ if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/runner_util)
 endif()
 
+if(EXECUTORCH_BUILD_EXTENSION_AOT_UTIL)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/aot_util)
+endif()
+
 option(EXECUTORCH_BUILD_XNNPACK "Build the backends/xnnpack directory" OFF)
 if(EXECUTORCH_BUILD_XNNPACK)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/xnnpack)
@@ -453,7 +461,6 @@ if(EXECUTORCH_BUILD_PYBIND)
       etdump
       executorch
       extension_data_loader
-      flatcc
       portable_ops_lib
       util
       torch

Package.swift

@@ -0,0 +1,65 @@
+// swift-tools-version:5.9
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import PackageDescription
+
+let url = "https://ossci-ios.s3.amazonaws.com/executorch"
+let version = "0.1.0"
+let coreml_sha256 = "1d2b8d2a5805a699eb39f347977894d3af5dfa763b298b926016550e9ffefda5"
+let executorch_sha256 = "39f19740a7c656d972e6082bae49583a6d4cc6396dea6ace2e4193688cef6225"
+let mps_sha256 = "866739b76baec70e603d331ff34ff9f028202fef69161f63a35d2e8a0cf502e9"
+let portable_sha256 = "6f761c0ae5651002e321bc6320604476ba0210f9383e535a2905cc1a74be55a3"
+let xnnpack_sha256 = "ef2cb2145a466a0a9e32489497c7f4880e4b582cea3883158b7ae24427d8ae7a"
+
+struct Framework {
+  let name: String
+  let checksum: String
+
+  func target() -> Target {
+    .binaryTarget(
+      name: name,
+      url: "\(url)/\(name)-\(version).zip",
+      checksum: checksum
+    )
+  }
+}
+
+let frameworks = [
+  Framework(
+    name: "coreml_backend",
+    checksum: coreml_sha256
+  ),
+  Framework(
+    name: "executorch",
+    checksum: executorch_sha256
+  ),
+  Framework(
+    name: "mps_backend",
+    checksum: mps_sha256
+  ),
+  Framework(
+    name: "portable_backend",
+    checksum: portable_sha256
+  ),
+  Framework(
+    name: "xnnpack_backend",
+    checksum: xnnpack_sha256
+  )
+]
+
+let package = Package(
+  name: "executorch",
+  platforms: [
+    .iOS(.v15),
+  ],
+  products: frameworks.map { framework in
+    .library(name: framework.name, targets: [framework.name])
+  },
+  targets: frameworks.map { $0.target() }
+)

backends/apple/coreml/partition/coreml_partitioner.py

@@ -17,6 +17,7 @@
     Partitioner,
     PartitionResult,
 )
+from executorch.exir.backend.utils import tag_constant_data
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
 from torch.fx.passes.operator_support import OperatorSupportBase
@@ -87,6 +88,8 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 node.meta["delegation_tag"] = tag
                 partition_tags[tag] = self.delegation_spec
 
+        tag_constant_data(exported_program)
+
         return PartitionResult(
             tagged_exported_program=exported_program, partition_tags=partition_tags
         )

backends/apple/coreml/runtime/delegate/ETCoreMLLogging.h

@@ -22,6 +22,7 @@ typedef NS_ERROR_ENUM(ETCoreMLErrorDomain, ETCoreMLError) {
     ETCoreMLErrorBrokenModel, // CoreML model doesn't match the input and output specification.
     ETCoreMLErrorCompilationFailed, // CoreML model failed to compile.
     ETCoreMLErrorModelCompilationNotSupported, // CoreML model compilation is not supported by the target.
+    ETCoreMLErrorModelProfilingNotSupported, // Model profiling is not supported by the target.
     ETCoreMLErrorModelSaveFailed, // Failed to save CoreML model to disk.
     ETCoreMLErrorModelCacheCreationFailed, // Failed to create model cache.
     ETCoreMLErrorInternalError, // Internal error.

backends/apple/coreml/runtime/inmemoryfs/inmemory_filesystem_py.cpp

@@ -14,6 +14,7 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/pytypes.h>
 #include <sstream>
+#include <stdexcept>
 #include <string>
 #include <sys/mman.h>
 #include <system_error>
@@ -151,8 +152,38 @@ pybind11::bytes flatten_directory_contents(const std::string& path) {
     return bytes == nullptr ? pybind11::none() : pybind11::reinterpret_steal<pybind11::object>((PyObject*)bytes);
 }
 
+/// Unflattens and writes the contents of the memory buffer at the specified path.
+///
+/// @param bytes  The bytes returned from `flatten_directory_contents`.
+/// @param path  The directory path
+bool unflatten_directory_contents(pybind11::bytes bytes, const std::string& path) {
+    using namespace inmemoryfs;
+
+    char* buffer = nullptr;
+    ssize_t length = 0;
+    if (PYBIND11_BYTES_AS_STRING_AND_SIZE(bytes.ptr(), &buffer, &length)) {
+        pybind11::pybind11_fail("Failed to extract contents of bytes object!");
+    }
+    std::shared_ptr<MemoryBuffer> memory_buffer =
+        MemoryBuffer::make_unowned((void*)buffer, static_cast<size_t>(length));
+    auto fs = inmemoryfs::make_from_buffer(memory_buffer);
+    if (!fs) {
+        pybind11::pybind11_fail("Failed to de-serialize bytes object!");
+        return false;
+    }
+    std::error_code ec;
+    std::filesystem::path fs_path(path);
+    auto canonical_path = std::filesystem::canonical(fs_path);
+    if (!fs->write_item_to_disk({}, canonical_path, true, ec)) {
+        pybind11::pybind11_fail("Failed to write the item to disk!");
+        return false;
+    }
+
+    return true;
+}
 } // namespace executorchcoreml
 
 PYBIND11_MODULE(executorchcoreml, mod) {
     mod.def("flatten_directory_contents", &executorchcoreml::flatten_directory_contents);
+    mod.def("unflatten_directory_contents", &executorchcoreml::unflatten_directory_contents);
 }