[ET-VK][examples] Create export script for Vulkan examples

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/vulkan/CMakeLists.txt

@@ -101,7 +101,7 @@ set_target_properties(vulkan_schema PROPERTIES LINKER_LANGUAGE CXX)
 target_include_directories(
   vulkan_schema
   INTERFACE
-    ${SCHEMA_INCLUDE_DIR}
+    $<BUILD_INTERFACE:${SCHEMA_INCLUDE_DIR}>
     $<BUILD_INTERFACE:${EXECUTORCH_ROOT}/third-party/flatbuffers/include>
 )
 

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -7,7 +7,7 @@
 # pyre-strict
 
 import logging
-from typing import Any, Callable, Dict, final, List, Mapping, Optional, Tuple
+from typing import Any, Callable, Dict, final, List, Mapping, Optional, Set, Tuple
 
 import executorch.backends.vulkan.utils as utils
 
@@ -17,6 +17,7 @@
     get_op_features,
     has_impl,
     OpFeatures,
+    OpKey,
     vulkan_supported_ops,
 )
 
@@ -55,11 +56,17 @@ def __init__(
         texture_limits: utils.ImageExtents,
         buffer_limit: int,
         require_dynamic_shape: bool = False,
+        operator_blocklist: Optional[Set[OpKey]] = None,
+        operator_allowlist: Optional[Set[OpKey]] = None,
     ) -> None:
         super().__init__()
         self.texture_limits: utils.ImageExtents = texture_limits
         self.buffer_limit = buffer_limit
         self.require_dynamic_shapes = require_dynamic_shape
+        self.operator_blocklist: Set[OpKey] = (
+            operator_blocklist if operator_blocklist is not None else set()
+        )
+        self.operator_allowlist = operator_allowlist
 
     def op_node_is_compatible(  # noqa: C901: Function is too complex
         self, node: torch.fx.Node, features: Optional[OpFeatures] = None
@@ -77,6 +84,17 @@ def op_node_is_compatible(  # noqa: C901: Function is too complex
             assert isinstance(first_arg, torch._ops.OpOverload)
             target = first_arg.name()
 
+        # Operator allow list is only used for torch ops
+        if (
+            utils.is_torch_op_node(node)
+            and (self.operator_allowlist is not None)
+            and (target not in self.operator_allowlist)
+        ):
+            return False, "op is not in allowlist"
+
+        if target in self.operator_blocklist:
+            return False, "op is in blocklist"
+
         # Extract the features for the node's operator, if no override was provided
         if features is None:
             if not has_impl(target):
@@ -93,7 +111,7 @@ def op_node_is_compatible(  # noqa: C901: Function is too complex
         if op_repsets.any_is_empty():
             return (
                 False,
-                "No valid representations for a tensor in the operation",
+                f"no valid representations for op {utils.node_io_str(node)}",
             )
 
         return True, "Op is compatible"
@@ -277,6 +295,8 @@ class VulkanPartitioner(Partitioner):
     def __init__(
         self,
         compile_options: Optional[Dict[str, Any]] = None,
+        operator_blocklist: Optional[List[OpKey]] = None,
+        operator_allowlist: Optional[List[OpKey]] = None,
     ) -> None:
         self.options: Dict[str, Any] = {}
         if compile_options is not None:
@@ -285,6 +305,18 @@ def __init__(
         compile_spec = parse_compile_options(self.options)
         self.delegation_spec = DelegationSpec(VulkanBackend.__name__, compile_spec)
 
+        self.operator_blocklist: Set[OpKey] = set()
+        if operator_blocklist is not None:
+            for entry in operator_blocklist or []:
+                self.operator_blocklist.add(entry)
+
+        self.operator_allowlist: Optional[Set[OpKey]] = None
+        if operator_allowlist is not None:
+            self.operator_allowlist = set()
+            for entry in operator_allowlist:
+                assert self.operator_allowlist is not None
+                self.operator_allowlist.add(entry)
+
     def ops_to_not_decompose(
         self, ep: ExportedProgram
     ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.Node], bool]]]:
@@ -308,6 +340,8 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult:
                 texture_limits,
                 buffer_limit,
                 require_dynamic_shape=self.options.get("require_dynamic_shapes", False),
+                operator_blocklist=self.operator_blocklist,
+                operator_allowlist=self.operator_allowlist,
             ),
             allows_single_node_partition=True,
         )

backends/vulkan/runtime/gen_vulkan_spv.py

@@ -1083,7 +1083,6 @@ def compile_spirv(shader_paths_pair) -> Tuple[str, str]:
             for spv_out_path, glsl_out_path in pool.map(
                 compile_spirv, self.output_file_map.items()
             ):
-                print(spv_to_glsl_map)
                 spv_to_glsl_map[spv_out_path] = glsl_out_path
 
         return spv_to_glsl_map

backends/vulkan/test/scripts/test_model.sh

@@ -0,0 +1,180 @@
+#!/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 -exu
+
+# Initialize variables
+RUN_BUILD=false
+RUN_CORRECTNESS_TEST=false
+RUN_CLEAN=false
+RUN_RECOMPILE=false
+MODEL_NAME=""
+OUTPUT_DIRECTORY="."
+
+# Parse arguments
+SKIP_NEXT=false
+for i in $(seq 1 $#); do
+  if [[ "$SKIP_NEXT" == true ]]; then
+    SKIP_NEXT=false
+    continue
+  fi
+
+  arg="${!i}"
+  case $arg in
+    --build|-b)
+      RUN_BUILD=true
+      ;;
+    --clean|-c)
+      RUN_CLEAN=true
+      ;;
+    --recompile|-rc)
+      RUN_RECOMPILE=true
+      ;;
+    --output_directory|-o)
+      next_i=$((i + 1))
+      if [[ $next_i -le $# ]]; then
+        OUTPUT_DIRECTORY="${!next_i}"
+        SKIP_NEXT=true
+      else
+        echo "Error: --output_directory|-o requires a value"
+        exit 1
+      fi
+      ;;
+    --*|-*)
+      echo "Unknown argument: $arg"
+      exit 1
+      ;;
+    *)
+      if [[ -z "$MODEL_NAME" ]]; then
+        MODEL_NAME="$arg"
+      else
+        echo "Multiple model names provided: $MODEL_NAME and $arg"
+        exit 1
+      fi
+      ;;
+  esac
+done
+
+# Determine execution mode based on parsed arguments
+if [[ "$RUN_BUILD" == true ]] && [[ -z "$MODEL_NAME" ]]; then
+  # Build-only mode
+  RUN_CORRECTNESS_TEST=false
+elif [[ "$RUN_BUILD" == true ]] && [[ -n "$MODEL_NAME" ]]; then
+  # Build and test mode
+  RUN_CORRECTNESS_TEST=true
+elif [[ "$RUN_BUILD" == false ]] && [[ -n "$MODEL_NAME" ]]; then
+  # Test-only mode
+  RUN_CORRECTNESS_TEST=true
+else
+  echo "Invalid argument combination. Usage:"
+  echo "  $0 --build|-b [--clean|-c] [--recompile|-rc] [-o|--output_directory DIR]                    # Build-only mode"
+  echo "  $0 model_name [--build|-b] [--clean|-c] [--recompile|-rc] [-o|--output_directory DIR]       # Test mode or build+test mode"
+  exit 1
+fi
+
+if [[ -z "${PYTHON_EXECUTABLE:-}" ]]; then
+  PYTHON_EXECUTABLE=python3
+fi
+which "${PYTHON_EXECUTABLE}"
+
+CMAKE_OUTPUT_DIR=cmake-out
+
+# Only set EXPORTED_MODEL if running correctness test
+if [[ "${RUN_CORRECTNESS_TEST}" == true ]]; then
+  EXPORTED_MODEL=${MODEL_NAME}_vulkan
+fi
+
+
+clean_build_directory() {
+  echo "Cleaning build directory: ${CMAKE_OUTPUT_DIR}"
+  rm -rf ${CMAKE_OUTPUT_DIR}
+}
+
+recompile() {
+  cmake --build cmake-out -j64 --target install
+}
+
+build_core_libraries_and_devtools() {
+  echo "Building core libraries and devtools with comprehensive Vulkan support..."
+
+  # Build core libraries with all required components
+  cmake . \
+    -DCMAKE_INSTALL_PREFIX=cmake-out \
+    -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \
+    -DEXECUTORCH_BUILD_KERNELS_CUSTOM_AOT=ON \
+    -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \
+    -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
+    -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
+    -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \
+    -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
+    -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON \
+    -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
+    -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=ON \
+    -DEXECUTORCH_BUILD_DEVTOOLS=ON \
+    -DEXECUTORCH_BUILD_VULKAN=ON \
+    -DEXECUTORCH_BUILD_XNNPACK=ON \
+    -Bcmake-out && \
+  cmake --build cmake-out -j64 --target install
+
+  # Build devtools example runner
+  cmake examples/devtools \
+    -DCMAKE_INSTALL_PREFIX=cmake-out \
+    -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \
+    -DEXECUTORCH_BUILD_VULKAN=ON \
+    -Bcmake-out/examples/devtools && \
+  cmake --build cmake-out/examples/devtools -j16 --config Release
+}
+
+run_example_runner() {
+  ./${CMAKE_OUTPUT_DIR}/examples/devtools/example_runner -bundled_program_path "${OUTPUT_DIRECTORY}/${EXPORTED_MODEL}.bpte" -output_verification
+}
+
+test_bundled_model_with_vulkan() {
+  # Export model as bundled program with Vulkan backend
+  "${PYTHON_EXECUTABLE}" -m examples.vulkan.export --model_name="${MODEL_NAME}" --output_dir="${OUTPUT_DIRECTORY}" --bundled
+
+  # Update exported model name for bundled program
+  EXPORTED_MODEL="${MODEL_NAME}_vulkan"
+
+  # Verify the exported bundled model exists
+  if [[ ! -f "${OUTPUT_DIRECTORY}/${EXPORTED_MODEL}.bpte" ]]; then
+    echo "Error: Failed to export bundled model ${MODEL_NAME} with Vulkan backend"
+    exit 1
+  fi
+
+  # Note: Running bundled programs may require different executor runner
+  echo "Bundled program created successfully. Use appropriate bundled program runner to test."
+
+  run_example_runner
+}
+
+
+# Main execution
+if [[ "${RUN_BUILD}" == true ]]; then
+  if [[ "${RUN_CLEAN}" == true ]]; then
+    clean_build_directory
+  fi
+  build_core_libraries_and_devtools
+fi
+
+if [[ "${RUN_RECOMPILE}" == true ]]; then
+  recompile
+fi
+
+if [[ "${RUN_CORRECTNESS_TEST}" == true ]]; then
+  echo "Testing ${MODEL_NAME} with Vulkan backend..."
+  # Always use bundled program testing
+  test_bundled_model_with_vulkan
+
+  # Check if test completed successfully
+  if [[ $? -eq 0 ]]; then
+    echo "Vulkan model test completed successfully!"
+  else
+    echo "Vulkan model test failed!"
+    exit 1
+  fi
+fi