[ET-VK][CI] Add vulkan CI to test exporting and running models (pytorch#13295)

Summary:
Title says it all!

cc @manuelcandales @cbilgin

---------

Co-authored-by: ssjia <[email protected]>

Author: 2 people

.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

@@ -864,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

backends/vulkan/cmake/ShaderLibrary.cmake

@@ -49,6 +49,15 @@ function(gen_vulkan_shader_lib_cpp shaders_path)
   set(VULKAN_SHADERGEN_ENV "")
   set(VULKAN_SHADERGEN_OUT_PATH ${CMAKE_BINARY_DIR}/vulkan_compute_shaders)
 
+  set(GEN_SPV_ARGS "--optimize")
+  if(DEFINED ENV{ETVK_USING_SWIFTSHADER})
+    if("$ENV{ETVK_USING_SWIFTSHADER}" STREQUAL "1"
+       OR "$ENV{ETVK_USING_SWIFTSHADER}" STREQUAL "True"
+    )
+      list(APPEND GEN_SPV_ARGS "--replace-u16vecn")
+    endif()
+  endif()
+
   add_custom_command(
     COMMENT "Generating Vulkan Compute Shaders"
     OUTPUT ${VULKAN_SHADERGEN_OUT_PATH}/spv.cpp
@@ -58,7 +67,7 @@ function(gen_vulkan_shader_lib_cpp shaders_path)
       ${shaders_path} --output-path ${VULKAN_SHADERGEN_OUT_PATH}
       --glslc-path=${GLSLC_PATH}
       --tmp-dir-path=${VULKAN_SHADERGEN_OUT_PATH}/shader_cache/ --env
-      ${VULKAN_GEN_ARG_ENV} --optimize
+      ${VULKAN_GEN_ARG_ENV} ${GEN_SPV_ARGS}
     DEPENDS ${shaders_path}/*
             ${EXECUTORCH_ROOT}/backends/vulkan/runtime/gen_vulkan_spv.py
   )

backends/vulkan/test/utils.py

@@ -6,6 +6,7 @@
 
 
 import logging
+from collections import OrderedDict
 from typing import List, Optional, Tuple
 
 import executorch.backends.vulkan.utils as utils
@@ -114,6 +115,10 @@ def check_outputs_equal(
     Helper function that checks if model output and reference output are equal with some tolerance.
     Returns True if equal, False otherwise.
     """
+    # Convert OrderedDict to list if needed
+    if isinstance(ref_output, OrderedDict):
+        ref_output = list(ref_output.values())
+
     # Compare the result from executor and eager mode directly
     if isinstance(ref_output, tuple) or isinstance(ref_output, list):
         # Multiple outputs executor always returns tuple, even if there is one output