Update on "[ET-VK][ez] Enable dynamic shape support when using push constants"

## Changes

* Call `encode_execute()` upon resize in `VulkanBackend.cpp`
* Minor update to `DispatchNode` to store push constant data array as a persistent member of the class

## Motivation

Passing in tensor metadata (i.e. sizes, strides) via push constants is typically more performant than passing them via a UBO (uniform buffer object). However, currently dynamic shapes do not work when push constants are used as I realized that the tensor metadata contained in the push constants do not get updated.

It appears that that `vkCmdPushConstants` sets the push constants when encoding the command buffer, however the push constants will not be updated if the command buffer is submitted for execution multiple times. Therefore, to update push constant values **the command buffer needs to be re-encoded**.

## Performance Impact

This may add a small performance overhead (i.e. re-encoding the command buffer) when executing models with dynamic shapes. Models that do not trigger tensor resizing will not be impacted.

However, I measured the impact on a llama 3.2 1B model and the impact of re-encoding a command buffer appears to be negligible.

In any case, re-encoding the command buffer is a "necessary evil" when working with dynamic shapes, otherwise the tensor metadata seen by shaders may never get updated. Furthermore, re-encoding the command buffer can allow an opportunity to adjust global work group sizing to match current tensor sizes, which may have a huge performance impact when maximum tensor sizes far exceeds what tensor sizes will realistically be during inference (one instance of this is for transformer models when the max sequence length is very long).

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

[ghstack-poisoned]

Author: SS-JIA

.Package.swift/kernels_portable/dummy.swift

@@ -42,10 +42,6 @@ say "Installing CoreML Backend Requirements"
 
 ./backends/apple/coreml/scripts/install_requirements.sh
 
-say "Installing MPS Backend Requirements"
-
-./backends/apple/mps/install_requirements.sh
-
 say "Exporting Models"
 
 python3 -m examples.portable.scripts.export --model_name="$MODEL_NAME" --segment_alignment=0x4000

.Package.swift/kernels_portable_debug/dummy.swift

@@ -188,11 +188,6 @@ jobs:
             backends/apple/coreml/scripts/install_requirements.sh
         fi
 
-        if [[ ${{ matrix.config }} == *"mps"* ]]; then
-          PYTHON_EXECUTABLE=python ${CONDA_RUN} --no-capture-output \
-            backends/apple/mps/install_requirements.sh
-        fi
-
         # Install requirements for export_llama
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash examples/models/llama/install_requirements.sh
 
@@ -379,10 +374,6 @@ jobs:
         # Install CoreML Backend Requirements
         PYTHON_EXECUTABLE=python ${CONDA_RUN} --no-capture-output \
           backends/apple/coreml/scripts/install_requirements.sh
-
-        # Install MPS Backend Requirements
-        PYTHON_EXECUTABLE=python ${CONDA_RUN} --no-capture-output \
-          backends/apple/mps/install_requirements.sh
         echo "::endgroup::"
 
         echo "::group::Build ExecuTorch iOS frameworks"

.ci/scripts/test_ios_ci.sh

@@ -154,7 +154,6 @@ jobs:
           "backend_xnnpack"
           "kernels_custom"
           "kernels_optimized"
-          "kernels_portable"
           "kernels_quantized"
           "threadpool"
         )
@@ -169,10 +168,6 @@ jobs:
         PYTHON_EXECUTABLE=python ${CONDA_RUN} --no-capture-output \
         backends/apple/coreml/scripts/install_requirements.sh
 
-        # Install MPS Backend Requirements
-        PYTHON_EXECUTABLE=python ${CONDA_RUN} --no-capture-output \
-        backends/apple/mps/install_requirements.sh
-
         # Build iOS Frameworks
         PYTHON_EXECUTABLE=python ${CONDA_RUN} --no-capture-output scripts/build_apple_frameworks.sh
 
@@ -307,10 +302,6 @@ jobs:
         # Install CoreML Backend Requirements
         PYTHON_EXECUTABLE=python ${CONDA_RUN} --no-capture-output \
           backends/apple/coreml/scripts/install_requirements.sh
-
-        # Install MPS Backend Requirements
-        PYTHON_EXECUTABLE=python ${CONDA_RUN} --no-capture-output \
-          backends/apple/mps/install_requirements.sh
         echo "::endgroup::"
 
         echo "::group::Build ExecuTorch iOS frameworks"

.github/workflows/apple-perf.yml

@@ -414,11 +414,7 @@ jobs:
         # Setup executorch
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/setup-macos.sh --build-tool cmake
 
-        if [[ "${MODE}" == "mps" ]]; then
-          # Install mps delegate
-          PYTHON_EXECUTABLE=python ${CONDA_RUN} bash backends/apple/mps/install_requirements.sh
-          echo "Finishing installing mps."
-        elif [[ "${MODE}" == "coreml" ]]; then
+        if [[ "${MODE}" == "coreml" ]]; then
           # Install coreml delegate
           PYTHON_EXECUTABLE=python ${CONDA_RUN} bash backends/apple/coreml/scripts/install_requirements.sh
           echo "Finishing installing coreml."
@@ -504,8 +500,6 @@ jobs:
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/setup-macos.sh --build-tool "${BUILD_TOOL}"
         PYTHON_EXECUTABLE=python ${CONDA_RUN} bash backends/apple/coreml/scripts/install_requirements.sh
         echo "Finishing installing coreml."
-        PYTHON_EXECUTABLE=python ${CONDA_RUN} bash backends/apple/mps/install_requirements.sh
-        echo "Finishing installing mps."
 
         # Build and test coreml model
         MODELS=(mv3 ic4 resnet50 edsr mobilebert w2l)

.github/workflows/apple.yml

@@ -39,13 +39,16 @@
 # ~~~
 # cmake-format -i CMakeLists.txt
 # ~~~
-# It should also be cmake-lint clean.
+# It should also be checked with a linter via
+# ~~~
+# cmake-lint CMakeLists.txt
+# ~~~
 #
 
 cmake_minimum_required(VERSION 3.24)
 project(executorch)
 
-# MARK: - Start EXECUTORCH_H12025_BUILD_MIGRATION --------------------------------------------------
+# MARK: - Start EXECUTORCH_H12025_BUILD_MIGRATION
 
 include(${PROJECT_SOURCE_DIR}/tools/cmake/common/preset.cmake)
 include(${PROJECT_SOURCE_DIR}/tools/cmake/Utils.cmake)
@@ -82,24 +85,25 @@ include(${PROJECT_SOURCE_DIR}/tools/cmake/preset/default.cmake)
 # Print all the configs that were called with announce_configured_options.
 print_configured_options()
 
-# MARK: - End EXECUTORCH_H12025_BUILD_MIGRATION ----------------------------------------------------
+# MARK: - End EXECUTORCH_H12025_BUILD_MIGRATION
 
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
-# Setup RPATH.
-# See https://gitlab.kitware.com/cmake/community/-/wikis/doc/cmake/RPATH-handling
+# Setup RPATH. See
+# https://gitlab.kitware.com/cmake/community/-/wikis/doc/cmake/RPATH-handling
 # Use separate rpaths during build and install phases
 set(CMAKE_SKIP_BUILD_RPATH OFF)
 # Don't use the install-rpath during the build phase
 set(CMAKE_BUILD_WITH_INSTALL_RPATH ON)
 # Automatically add all linked folders that are NOT in the build directory to
 # the rpath (per library?)
-# TODO: Doesn't work for us right now because we are not installing .so's into the
-# correct locations. For example we have libcustom_ops_aot_lib.so depending on
-# _portable_lib.so, which was eventually put under <site-packages>/executorch/extension/pybindings/
-# but this rpath is not automatically added because at build time it seems `portable_lib`
-# is being built under the same directory, so no extra rpath is being added. To
-# properly fix this we need to install `portable_lib` into the correct path.
+# TODO: Doesn't work for us right now because we are
+# not installing .so's into the correct locations. For example we have
+# libcustom_ops_aot_lib.so depending on _portable_lib.so, which was eventually
+# put under <site-packages>/executorch/extension/pybindings/ but this rpath is
+# not automatically added because at build time it seems `portable_lib` is being
+# built under the same directory, so no extra rpath is being added. To properly
+# fix this we need to install `portable_lib` into the correct path.
 set(CMAKE_INSTALL_RPATH_USE_LINK_PATH ON)
 # ------------------------------ OPTIONS -------------------------------------
 # WARNING: Please don't add example specific options in this CMakeLists.txt.
@@ -143,6 +147,11 @@ else()
   set(CMAKE_CXX_FLAGS_RELEASE "-O2 ${CMAKE_CXX_FLAGS_RELEASE}")
 endif()
 
+if(EXECUTORCH_BUILD_TESTS)
+  set(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON)
+  include(CTest)
+endif()
+
 add_subdirectory(third-party)
 
 if(EXECUTORCH_BUILD_EXTENSION_TRAINING)
@@ -172,7 +181,7 @@ endif()
 
 if(NOT DEFINED FXDIV_SOURCE_DIR)
   set(ORIGINAL_CMAKE_POSITION_INDEPENDENT_CODE_FLAG
-    ${CMAKE_POSITION_INDEPENDENT_CODE}
+      ${CMAKE_POSITION_INDEPENDENT_CODE}
   )
   set(FXDIV_SOURCE_DIR "backends/xnnpack/third-party/FXdiv")
   add_subdirectory("${FXDIV_SOURCE_DIR}")
@@ -254,11 +263,6 @@ if(EXECUTORCH_BUILD_PTHREADPOOL)
   )
 endif()
 
-if(EXECUTORCH_BUILD_TESTS)
-  set(EXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR ON)
-  include(CTest)
-endif()
-
 # TODO(dbort): Fix these warnings and remove this flag.
 set(_common_compile_options -Wno-deprecated-declarations -fPIC)
 
@@ -276,7 +280,10 @@ if(NOT "${_repo_dir_name}" STREQUAL "executorch")
       "fix for this restriction."
   )
 endif()
-set(_common_include_directories ${CMAKE_CURRENT_SOURCE_DIR}/.. ${CMAKE_CURRENT_SOURCE_DIR}/runtime/core/portable_type/c10)
+set(_common_include_directories
+    ${CMAKE_CURRENT_SOURCE_DIR}/..
+    ${CMAKE_CURRENT_SOURCE_DIR}/runtime/core/portable_type/c10
+)
 
 #
 # The `_<target>_srcs` lists are defined by including ${EXECUTORCH_SRCS_FILE}.
@@ -310,9 +317,9 @@ endif()
 # Detect if an Android toolchain is set.
 if(CMAKE_TOOLCHAIN_FILE MATCHES ".*android\.toolchain\.cmake$")
   set(CMAKE_TOOLCHAIN_ANDROID ON)
-if(NOT ANDROID_PLATFORM)
-  set(ANDROID_PLATFORM android-30)
-endif()
+  if(NOT ANDROID_PLATFORM)
+    set(ANDROID_PLATFORM android-30)
+  endif()
 else()
   set(CMAKE_TOOLCHAIN_ANDROID OFF)
 endif()
@@ -334,7 +341,6 @@ if(EXECUTORCH_USE_CPP_CODE_COVERAGE)
   endif()
 endif()
 
-
 #
 # program_schema: Generated .h files from schema/*.fbs inputs
 #
@@ -376,7 +382,9 @@ endif()
 target_include_directories(
   executorch_core PUBLIC ${_common_include_directories}
 )
-target_compile_definitions(executorch_core PUBLIC C10_USING_CUSTOM_GENERATED_MACROS)
+target_compile_definitions(
+  executorch_core PUBLIC C10_USING_CUSTOM_GENERATED_MACROS
+)
 target_compile_options(executorch_core PUBLIC ${_common_compile_options})
 if(MAX_KERNEL_NUM)
   target_compile_definitions(
@@ -386,9 +394,7 @@ endif()
 
 if(EXECUTORCH_BUILD_PYBIND AND APPLE)
   # shared version
-  add_library(
-    executorch_core_shared SHARED ${_executorch_core__srcs}
-  )
+  add_library(executorch_core_shared SHARED ${_executorch_core__srcs})
   target_link_libraries(executorch_core_shared PRIVATE program_schema)
   if(DL_LIBRARY_EXISTS)
     # For dladdr()
@@ -397,7 +403,9 @@ if(EXECUTORCH_BUILD_PYBIND AND APPLE)
   target_include_directories(
     executorch_core_shared PUBLIC ${_common_include_directories}
   )
-  target_compile_definitions(executorch_core_shared PUBLIC C10_USING_CUSTOM_GENERATED_MACROS)
+  target_compile_definitions(
+    executorch_core_shared PUBLIC C10_USING_CUSTOM_GENERATED_MACROS
+  )
   target_compile_options(
     executorch_core_shared PUBLIC ${_common_compile_options}
   )
@@ -430,9 +438,8 @@ target_link_options_shared_lib(executorch)
 # operators necessary for the models that will run.
 #
 if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
-  # find pytorch lib here to make it available to all
-  # sub-directories. Find it before including portable so that
-  # optimized_portable_kernels can use it.
+  # find pytorch lib here to make it available to all sub-directories. Find it
+  # before including portable so that optimized_portable_kernels can use it.
   find_package_torch_headers()
 endif()
 
@@ -458,24 +465,50 @@ endif()
 
 # Install `executorch` library as well as `executorch-config.cmake` under
 # ${CMAKE_INSTALL_PREFIX}/
-install(DIRECTORY runtime/core/  DESTINATION include/executorch/runtime/core FILES_MATCHING PATTERN "*.h")
-install(DIRECTORY runtime/kernel/  DESTINATION include/executorch/runtime/kernel FILES_MATCHING PATTERN "*.h")
-install(DIRECTORY runtime/platform/  DESTINATION include/executorch/runtime/platform FILES_MATCHING PATTERN "*.h")
-install(DIRECTORY extension/kernel_util/  DESTINATION include/executorch/extension/kernel_util FILES_MATCHING PATTERN "*.h")
-install(DIRECTORY extension/tensor/  DESTINATION include/executorch/extension/tensor FILES_MATCHING PATTERN "*.h")
-install(DIRECTORY extension/threadpool/  DESTINATION include/executorch/extension/threadpool FILES_MATCHING PATTERN "*.h")
+install(
+  DIRECTORY runtime/core/
+  DESTINATION include/executorch/runtime/core
+  FILES_MATCHING
+  PATTERN "*.h"
+)
+install(
+  DIRECTORY runtime/kernel/
+  DESTINATION include/executorch/runtime/kernel
+  FILES_MATCHING
+  PATTERN "*.h"
+)
+install(
+  DIRECTORY runtime/platform/
+  DESTINATION include/executorch/runtime/platform
+  FILES_MATCHING
+  PATTERN "*.h"
+)
+install(
+  DIRECTORY extension/kernel_util/
+  DESTINATION include/executorch/extension/kernel_util
+  FILES_MATCHING
+  PATTERN "*.h"
+)
+install(
+  DIRECTORY extension/tensor/
+  DESTINATION include/executorch/extension/tensor
+  FILES_MATCHING
+  PATTERN "*.h"
+)
+install(
+  DIRECTORY extension/threadpool/
+  DESTINATION include/executorch/extension/threadpool
+  FILES_MATCHING
+  PATTERN "*.h"
+)
 install(
   TARGETS executorch executorch_core
-  DESTINATION lib
   INCLUDES
   DESTINATION ${_common_include_directories}
 )
-install(FILES tools/cmake/executorch-config.cmake DESTINATION lib/cmake/ExecuTorch)
-
-# Add googletest if any test targets should be built
-if(BUILD_TESTING)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third-party/googletest)
-endif()
+install(FILES tools/cmake/executorch-config.cmake
+        DESTINATION lib/cmake/ExecuTorch
+)
 
 if(EXECUTORCH_BUILD_ARM_BAREMETAL)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/arm)
@@ -613,17 +646,14 @@ if(EXECUTORCH_BUILD_PYBIND)
   endif()
 
   if(EXECUTORCH_BUILD_XNNPACK)
-    # need to explicitly specify XNNPACK and microkernels-prod
-    # here otherwise uses XNNPACK and microkernel-prod symbols from libtorch_cpu
+    # need to explicitly specify XNNPACK and microkernels-prod here otherwise
+    # uses XNNPACK and microkernel-prod symbols from libtorch_cpu
     list(APPEND _dep_libs xnnpack_backend XNNPACK microkernels-prod)
   endif()
 
   # compile options for pybind
-  set(_pybind_compile_options
-      -Wno-deprecated-declarations
-      -fPIC
-      -frtti
-      -fexceptions
+  set(_pybind_compile_options -Wno-deprecated-declarations -fPIC -frtti
+                              -fexceptions
   )
 
   # util lib