Update on "[ET-VK] Add mechanism to trigger command buffer re-encode only when necessary"

## Context

Dynamic shape models currently will require the command buffer to be re-encoded every inference. However, this introduces a significant overhead when running models that require dynamic shapes.

The reality is that a command buffer  re-encode may not be needed every frame. A command buffer re-encode will only be needed when:

1. Shader dispatch parameters change; i.e. new tensor sizes require a completely different compute shader, require new local work group sizing, or require new work group grid size (i.e. global work group size / local work group size)
2. Push constants containing tensor metadata need to be updated

This diff aims to reduce the overhead of triggering tensor shape change by detecting when a command buffer re-encode is actually needed.

## Changes

`ComputeGraph`:
* Introduce `requires_reencode` flag to `ComputeGraph` to indicate when a command buffer re-encode is needed.
* Introduce a `std::set<ValueRef>` tracking which values were updated when propagating tensor sizes
  * "update" can be one of two things: 1) tensor sizes changed 2) symint value changed

`DispatchNode`:
* When propagating new tensor sizes, only execute the resize function if any of the values participating in the computation have been updated
* Mark `requries_reencode` if any push constants associated with tensor metadata need to be udpated

`DynamicDispatchNode`:
* Only recompute compute shader dispatch params if any of the values participating in the computation have been updated
* Mark `requires_reencode` if 1) a new compute shader is required, 2) local work group size changed, 3) work group grid size changed

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

[ghstack-poisoned]

Author: ssjia

.ci/scripts/setup-vulkan-linux-deps.sh

@@ -23,6 +23,7 @@ install_swiftshader() {
 
   export VK_ICD_FILENAMES="${_swiftshader_dir}/swiftshader/build/Linux/vk_swiftshader_icd.json"
   export LD_LIBRARY_PATH="${_swiftshader_dir}/swiftshader/build/Linux/"
+  export ETVK_USING_SWIFTSHADER=1
 }
 
 install_vulkan_sdk() {

.github/workflows/pull.yml

@@ -801,6 +801,8 @@ jobs:
       id-token: write
       contents: read
     strategy:
+      matrix:
+        enable-etdump: ['', '--enable-etdump']
       fail-fast: false
     with:
       runner: linux.2xlarge
@@ -820,7 +822,7 @@ jobs:
         source .ci/scripts/setup-emscripten.sh
 
         # Test selective build
-        bash scripts/build_wasm_tests.sh
+        bash scripts/build_wasm_tests.sh ${{ matrix.enable-etdump }}
 
         # Install Jest
         cd cmake-out-wasm/extension/wasm/test
@@ -862,6 +864,43 @@ jobs:
         PYTHON_EXECUTABLE=python bash examples/nxp/run_aot_example.sh
 
 
+  test-vulkan-models-linux:
+    name: test-vulkan-models-linux
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    permissions:
+      id-token: write
+      contents: read
+    with:
+      runner: linux.2xlarge
+      docker-image: ci-image:executorch-ubuntu-22.04-clang12
+      submodules: 'recursive'
+      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}"
+
+        # Setup swiftshader and Vulkan SDK which are required to build the Vulkan delegate
+        source .ci/scripts/setup-vulkan-linux-deps.sh
+
+        # Setup python
+        PYTHON_EXECUTABLE=python \
+        CMAKE_ARGS="-DEXECUTORCH_BUILD_VULKAN=ON" \
+        .ci/scripts/setup-linux.sh --build-tool "cmake"
+
+        PYTHON_EXECUTABLE=python bash backends/vulkan/test/scripts/test_model.sh --build
+
+        # Test models serially
+        models="mv2 mv3 edsr resnet18 resnet50 dl3"
+        for model in $models; do
+          python -m examples.vulkan.export --model_name=$model --test
+        done
+
+
+
   nxp-build-test:
     name: nxp-build-test
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main

.gitignore

@@ -24,6 +24,7 @@ pip-out/
 # Any exported models and profiling outputs
 *.bin
 *.model
+*.etdump
 tokenizer.json
 *.pte
 *.ptd
@@ -58,6 +59,7 @@ xcuserdata/
 /include/
 /share/
 /version.py
+*.csv
 
 # Android
 *.aar

CMakeLists.txt

@@ -763,6 +763,10 @@ if(EXECUTORCH_BUILD_PYBIND)
     list(APPEND _dep_libs xnnpack_backend XNNPACK xnnpack-microkernels-prod)
   endif()
 
+  if(EXECUTORCH_BUILD_VULKAN)
+    list(APPEND _dep_libs vulkan_backend)
+  endif()
+
   # compile options for pybind
   set(_pybind_compile_options -Wno-deprecated-declarations -fPIC -frtti
                               -fexceptions

backends/apple/coreml/TARGETS

@@ -120,6 +120,7 @@ runtime.python_test(
         "test/*.py",
     ]),
     deps = [
+        "fbsource//third-party/pypi/coremltools:coremltools",
         "fbsource//third-party/pypi/pytest:pytest",
         ":partitioner",
         ":quantizer",