Update base for Update on "[ET-VK][Llama] Apply XNNPACK partitoner as well when lowering to Vulkan"

## Context

The final logit linear layer in the Transformer architecture has extremely large tensors, since the output and weight tensors will have a tensor with dim equal to the vocabulary size, which may be extremely large. Because of this, image textures cannot be used to execute the op when running with the Vulkan delegate, so an implementation using buffer based tensors must be used.

Unfortunately, Vulkan does not have a performant implementation of linear with buffer based tensors at the moment. As a result, if this final linear layer is executed in Vulkan, model inference is extremely slow.

## Changes

The below diff will prevent the final logit linear layer from being delegated to Vulkan by enforcing a GPU buffer limit.

This diff modifies the export llama script to apply the XNNPACK partitioner after the Vulkan partitioner if lowering to Vulkan, to ensure that remaining ops will be accelerated with XNNPACK. 4 bit quantization will also apply an additional Quantizer after applying the Vulkan quantizer (which will skip the final logit linear layer) so that the final logit linear can be quantized as well.

## Long Term

This is a temporary measure while an optimized buffer based linear implementation is developed. Once the Vulkan implementation achieves parity with XNNPACK, the final logit linear will be delegated to Vulkan once more.

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

[ghstack-poisoned]

Author: SS-JIA

.ci/docker/ci_commit_pins/pytorch.txt

@@ -1 +1 @@
-c8a648d4dffb9f0133ff4a2ea0e660b42105d3ad
+19eff28ff3f19b50da46f5a9ff5f4d4d213806fe

.ci/scripts/gather_test_models.py

@@ -24,6 +24,7 @@
         "ic4": "linux.12xlarge",
         "resnet50": "linux.12xlarge",
         "llava": "linux.12xlarge",
+        "llama3_2_vision_encoder": "linux.12xlarge",
         # This one causes timeout on smaller runner, the root cause is unclear (T161064121)
         "dl3": "linux.12xlarge",
         "emformer_join": "linux.12xlarge",

.ci/scripts/test_model.sh

@@ -77,7 +77,7 @@ test_model() {
     # Install requirements for export_llama
     bash examples/models/llama/install_requirements.sh
     # Test export_llama script: python3 -m examples.models.llama.export_llama
-    "${PYTHON_EXECUTABLE}" -m examples.models.llama.export_llama -c examples/models/llama/params/demo_rand_params.pth -p examples/models/llama/params/demo_config.json
+    "${PYTHON_EXECUTABLE}" -m examples.models.llama.export_llama --model "${MODEL_NAME}" -c examples/models/llama/params/demo_rand_params.pth -p examples/models/llama/params/demo_config.json
     run_portable_executor_runner
     rm "./${MODEL_NAME}.pte"
   fi

.github/workflows/pull.yml

@@ -72,6 +72,10 @@ jobs:
         conda activate "${CONDA_ENV}"
 
         MODEL_NAME=${{ matrix.model }}
+        # Install requirements for llama vision
+        if [[ "$MODEL_NAME" == "llama3_2_vision_encoder" ]]; then
+          bash examples/models/llama3_2_vision/install_requirements.sh
+        fi
         BUILD_TOOL=${{ matrix.build-tool }}
         BACKEND=${{ matrix.backend }}
         DEMO_BACKEND_DELEGATION=${{ matrix.demo_backend_delegation }}

.github/workflows/trunk.yml

@@ -58,7 +58,11 @@ jobs:
         bash .ci/scripts/setup-conda.sh
         # Setup MacOS dependencies as there is no Docker support on MacOS atm
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/setup-macos.sh "${BUILD_TOOL}"
-        # Build and test xecutorch
+        # Install requirements for llama vision
+        if [[ "$MODEL_NAME" == "llama3_2_vision_encoder" ]]; then
+          ${CONDA_RUN} bash examples/models/llama3_2_vision/install_requirements.sh
+        fi
+        # Build and test executorch
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" "${DEMO_BACKEND_DELEGATION}"
 
   test-custom-ops-macos:

.lintrunner.toml

@@ -196,3 +196,91 @@ command = [
     '@{{PATHSFILE}}',
 ]
 is_formatter = true
+
+[[linter]]
+code = 'NOSTDINC'
+include_patterns = [
+    "**/*.c",
+    "**/*.cpp",
+    "**/*.h",
+    "**/*.hpp",
+]
+exclude_patterns = [
+    '**/devtools/**',
+    '**/test/**',
+    '**/testing_util/**',
+    '**/third-party/**',
+    'backends/**',
+    'devtools/**',
+    'examples/**',
+    'extension/**',
+    'kernels/optimized/**',
+    'scripts/**',
+    'third-party/**',
+    'util/**',
+]
+command = [
+    'python',
+    '-m',
+    'lintrunner_adapters',
+    'run',
+    'grep_linter',
+    '--pattern=([^\\S\r\n]*#include\s*<(deque|exception|forward_list|functional|list|map|multimap|multiset|priority_queue|queue|set|stack|string|unordered_map|unordered_multimap|unordered_multiset|unordered_set|vector)>)',
+    '--linter-name=NOSTDINC',
+    '--error-name=Standard C++ container include in core',
+    """--error-description=\
+        Standard library containers should not be included in ExecuTorch core \
+        because they may call malloc, which is not allowed in core. \
+    """,
+    '--',
+    '@{{PATHSFILE}}',
+]
+
+[[linter]]
+code = 'NOTORCHINC'
+include_patterns = [
+    "**/*.c",
+    "**/*.cpp",
+    "**/*.h",
+    "**/*.hpp",
+]
+exclude_patterns = [
+    '**/devtools/**',
+    '**/fb/**',
+    '**/test/**',
+    '**/tests/**',
+    '**/testing_util/**',
+    '**/third-party/**',
+    'backends/**',
+    'codegen/templates/RegisterDispatchKeyCustomOps.cpp',
+    'codegen/templates/RegisterSchema.cpp',
+    'devtools/**',
+    'examples/**',
+    'exir/verification/bindings.cpp',
+    'extension/**',
+    'kernels/optimized/**',
+    'runtime/core/exec_aten/**',
+    'runtime/executor/tensor_parser_aten.cpp',
+    'scripts/**',
+    'test/**',
+    'third-party/**',
+    'util/**',
+]
+command = [
+    'python',
+    '-m',
+    'lintrunner_adapters',
+    'run',
+    'grep_linter',
+    '--pattern=#include\s+[<"](aten/|ATen/|torch/)',
+    '--linter-name=NOTORCHINC',
+    '--error-name=ATen or torch include',
+    """--error-description=\
+        PyTorch includes in ExecuTorch core are prohibited to prevent \
+        accidentally breaking core's requirements; please make sure this \
+        header complies (e.g., no streams/malloc/syscalls) and then include \
+        a patch to update this linter.\
+    """,
+    '--',
+    '@{{PATHSFILE}}',
+]

CMakeLists.txt

@@ -324,6 +324,19 @@ message(STATUS "Using python executable '${PYTHON_EXECUTABLE}'")
 set(_common_compile_options -Wno-deprecated-declarations -fPIC)
 
 # Let files say "include <executorch/path/to/header.h>".
+# TODO(#6475): This requires/assumes that the repo lives in a directory named
+# exactly `executorch`. Check the assumption first. Remove this check once we
+# stop relying on the assumption.
+cmake_path(GET CMAKE_CURRENT_SOURCE_DIR FILENAME _repo_dir_name)
+if(NOT "${_repo_dir_name}" STREQUAL "executorch")
+  message(
+    FATAL_ERROR
+      "The ExecuTorch repo must be cloned into a directory named exactly "
+      "`executorch`; found `${_repo_dir_name}`. See "
+      "https://github.com/pytorch/executorch/issues/6475 for progress on a "
+      "fix for this restriction."
+  )
+endif()
 set(_common_include_directories ${CMAKE_CURRENT_SOURCE_DIR}/..)
 
 #

backends/arm/TARGETS

@@ -26,12 +26,27 @@ python_library(
         "fbsource//third-party/serialization_lib/python/serializer:serializer",
         "fbsource//third-party/serialization_lib/python/tosa:tosa",
         ":arm_vela",
+        ":process_node",
         "//executorch/backends/arm/operators:lib",
         "//executorch/backends/arm/operators:node_visitor",
         "//executorch/backends/arm/_passes:passes",
     ],
 )
 
+python_library(
+    name = "process_node",
+    srcs = ["process_node.py"],
+    typing = True,
+    deps = [
+        "fbsource//third-party/serialization_lib/python/tosa:tosa",
+        "//executorch/backends/arm/operators:node_visitor",
+        "//executorch/backends/arm:tosa_mapping",
+        "//executorch/backends/arm:tosa_quant_utils",
+        "//executorch/backends/arm:tosa_utils",
+        "//executorch/exir:lib",
+    ],
+)
+
 python_library(
     name = "arm_vela",
     srcs = [

backends/arm/arm_backend.py

@@ -18,18 +18,17 @@
 import serializer.tosa_serializer as ts
 from executorch.backends.arm.arm_vela import vela_compile
 from executorch.backends.arm.operators.node_visitor import get_node_visitors
-from executorch.backends.arm.operators.op_output import process_output
-from executorch.backends.arm.operators.op_placeholder import process_placeholder
 
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.backends.arm._passes.arm_pass_manager import (
     ArmPassManager,
 )  # usort: skip
-from executorch.backends.arm.tosa_utils import (
-    dbg_fail,
-    dbg_tosa_dump,
+from executorch.backends.arm.process_node import (
     process_call_function,
+    process_output,
+    process_placeholder,
 )
+from executorch.backends.arm.tosa_utils import dbg_fail, dbg_tosa_dump
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from torch.export.exported_program import ExportedProgram