Fix OSS build for Vulkan Delegate (#2434)

Summary:
Pull Request resolved: #2434

## Context

Recent changes to the Vulkan backend and PyTorch Vulkan API requires that the CMake build for the Vulkan delegate be reworked. This PR updates the CMake rules for the Vulkan delegate and makes sure that the Vulkan delegate can build correctly.

imported-using-ghimport

Test Plan: Imported from OSS

Reviewed By: kirklandsign

Differential Revision: D54914149

Pulled By: SS-JIA

fbshipit-source-id: d613b40416dc31e638a1020c68ec81de3d3c2d15

Author: SS-JIA

backends/vulkan/CMakeLists.txt

@@ -8,7 +8,7 @@
 #
 # This file should be formatted with
 # ~~~
-# cmake-format --first-comment-is-literal=True CMakeLists.txt
+# cmake-format --first-comment-is-literal=True -i CMakeLists.txt
 # ~~~
 # It should also be cmake-lint clean.
 #
@@ -32,7 +32,10 @@ if(NOT FLATC_EXECUTABLE)
   set(FLATC_EXECUTABLE flatc)
 endif()
 
-# Include this file to access target_link_options_shared_lib
+# Include this file to access target_link_options_shared_lib This is required to
+# provide access to target_link_options_shared_lib which allows libraries to be
+# linked with the --whole-archive flag. This is required for libraries that
+# perform dynamic registration via static initialization.
 include(${EXECUTORCH_ROOT}/build/Utils.cmake)
 
 # ATen Vulkan Libs
@@ -48,12 +51,19 @@ set(COMMON_INCLUDES ${VULKAN_API_HEADERS} ${EXECUTORCH_ROOT}/..)
 file(GLOB_RECURSE vulkan_graph_cpp ${RUNTIME_PATH}/graph/*)
 
 add_library(vulkan_graph_lib STATIC ${vulkan_graph_cpp})
-
 target_include_directories(vulkan_graph_lib PRIVATE ${COMMON_INCLUDES})
+target_link_libraries(${LIBRARY_NAME} vulkan_api_lib)
+target_compile_options(vulkan_graph_lib PRIVATE ${VULKAN_CXX_FLAGS})
+# Link this library with --whole-archive due to dynamic operator registrations
+target_link_options_shared_lib(vulkan_graph_lib)
 
-target_link_libraries(vulkan_graph_lib vulkan_shader_lib)
+# Due to dynamic registrations, these libraries must be explicitly linked
+set(VULKAN_STANDARD_OPS_LIBS vulkan_graph_lib vulkan_graph_shaderlib)
 
-target_compile_options(vulkan_graph_lib PRIVATE ${VULKAN_CXX_FLAGS})
+# vulkan_graph_shaderlib
+
+set(VULKAN_GRAPH_SHADERS_PATH ${RUNTIME_PATH}/graph/ops/glsl/)
+vulkan_shader_library(${VULKAN_GRAPH_SHADERS_PATH} vulkan_graph_shaderlib)
 
 # Generate Files from flatc
 
@@ -85,48 +95,51 @@ target_include_directories(
 file(GLOB vulkan_backend_cpp ${RUNTIME_PATH}/*.cpp)
 
 add_library(vulkan_backend ${vulkan_backend_cpp})
-
-target_include_directories(vulkan_backend PRIVATE ${SCHEMA_INCLUDE_DIR})
-target_include_directories(vulkan_backend PRIVATE ${COMMON_INCLUDES})
-
-target_link_libraries(vulkan_backend PRIVATE vulkan_graph_lib)
-target_link_libraries(vulkan_backend PRIVATE vulkan_schema)
-target_link_libraries(vulkan_backend PRIVATE executorch)
-
+target_include_directories(vulkan_backend PRIVATE ${SCHEMA_INCLUDE_DIR}
+                                                  ${COMMON_INCLUDES})
+target_link_libraries(vulkan_backend PRIVATE vulkan_graph_lib vulkan_schema
+                                             executorch)
 target_compile_options(vulkan_backend PRIVATE ${VULKAN_CXX_FLAGS})
-
-# This is required to ensure that vulkan_backend gets linked with
-# --whole-archive since backends are registered via static variables that would
-# otherwise be discarded
+# Link this library with --whole-archive due to dynamic backend registration
 target_link_options_shared_lib(vulkan_backend)
 
 # Executor Runner
 
 if(NOT CMAKE_TOOLCHAIN_FILE MATCHES ".*iOS\.cmake$")
   set(VULKAN_RUNNER_SRCS ${_executor_runner__srcs})
   list(TRANSFORM VULKAN_RUNNER_SRCS PREPEND "${EXECUTORCH_ROOT}/")
+
   add_executable(vulkan_executor_runner ${VULKAN_RUNNER_SRCS})
-  target_link_libraries(vulkan_executor_runner ${_executor_runner_libs})
-  target_link_libraries(vulkan_executor_runner vulkan_schema)
-  target_link_libraries(vulkan_executor_runner vulkan_backend)
+  target_link_libraries(
+    vulkan_executor_runner ${_executor_runner_libs} vulkan_schema
+    vulkan_backend ${VULKAN_STANDARD_OPS_LIBS})
   target_compile_options(vulkan_executor_runner PUBLIC ${VULKAN_CXX_FLAGS})
 
   add_library(vulkan_executor_runner_lib STATIC ${VULKAN_RUNNER_SRCS})
-  target_link_libraries(vulkan_executor_runner_lib ${_executor_runner_libs})
-  target_link_libraries(vulkan_executor_runner_lib vulkan_schema)
-  target_link_libraries(vulkan_executor_runner_lib vulkan_backend)
+  target_link_libraries(vulkan_executor_runner_lib ${_executor_runner_libs}
+                        vulkan_schema vulkan_backend)
   target_compile_options(vulkan_executor_runner_lib PUBLIC ${VULKAN_CXX_FLAGS})
 endif()
 
 # Test targets
 
 if(EXECUTORCH_BUILD_GTESTS)
+  set(TEST_UTILS_HEADERS ${CMAKE_CURRENT_SOURCE_DIR}/test/utils)
+  file(GLOB TEST_UTILS_CPP ${CMAKE_CURRENT_SOURCE_DIR}/test/utils/*.cpp)
+
+  set(TEST_SHADERS_PATH ${CMAKE_CURRENT_SOURCE_DIR}/test/glsl)
+  vulkan_shader_library(${TEST_SHADERS_PATH} test_shaderlib)
+
   # vulkan_compute_api_test
-  set(TEST_CPP ${CMAKE_CURRENT_SOURCE_DIR}/test/vulkan_compute_api_test.cpp)
-  add_executable(vulkan_compute_api_test ${TEST_CPP})
-  target_include_directories(vulkan_compute_api_test PRIVATE ${COMMON_INCLUDES})
-  target_link_libraries(vulkan_compute_api_test vulkan_api_lib)
-  target_link_libraries(vulkan_compute_api_test vulkan_graph_lib)
-  target_link_libraries(vulkan_compute_api_test gtest_main)
+  set(COMPUTE_API_TEST_CPP
+      ${CMAKE_CURRENT_SOURCE_DIR}/test/vulkan_compute_api_test.cpp)
+
+  add_executable(vulkan_compute_api_test ${COMPUTE_API_TEST_CPP}
+                                         ${TEST_UTILS_CPP})
+  target_include_directories(vulkan_compute_api_test
+                             PRIVATE ${COMMON_INCLUDES} ${TEST_UTILS_HEADERS})
+  target_link_libraries(
+    vulkan_compute_api_test PRIVATE gtest_main ${VULKAN_STANDARD_OPS_LIBS}
+                                    test_shaderlib)
   target_compile_options(vulkan_compute_api_test PRIVATE ${VULKAN_CXX_FLAGS})
 endif()

backends/vulkan/cmake/ATenVulkan.cmake

@@ -8,7 +8,7 @@
 #
 # This file should be formatted with
 # ~~~
-# cmake-format --first-comment-is-literal=True CMakeLists.txt
+# cmake-format --first-comment-is-literal=True -i ATenVulkan.cmake
 # ~~~
 # It should also be cmake-lint clean.
 #
@@ -22,20 +22,6 @@ if(NOT VULKAN_THIRD_PARTY_PATH)
   set(VULKAN_THIRD_PARTY_PATH ${CMAKE_CURRENT_SOURCE_DIR}/../third-party)
 endif()
 
-# Shader Codegen
-
-# Trigger Shader code generation
-set(USE_VULKAN ON)
-set(VULKAN_CODEGEN_CMAKE_PATH ${PYTORCH_PATH}/cmake/VulkanCodegen.cmake)
-if(NOT EXISTS ${VULKAN_CODEGEN_CMAKE_PATH})
-  message(
-    FATAL_ERROR
-      "Cannot perform SPIR-V codegen because " ${VULKAN_CODEGEN_CMAKE_PATH}
-      " does not exist. Please make sure that submodules are initialized"
-      " and updated.")
-endif()
-include(${PYTORCH_PATH}/cmake/VulkanCodegen.cmake)
-
 # Source paths and compile settings
 
 set(ATEN_PATH ${PYTORCH_PATH}/aten/src)
@@ -65,19 +51,56 @@ list(APPEND VULKAN_API_HEADERS ${VOLK_PATH})
 list(APPEND VULKAN_API_HEADERS ${VMA_PATH})
 
 target_include_directories(vulkan_api_lib PRIVATE ${VULKAN_API_HEADERS})
-
 target_compile_options(vulkan_api_lib PRIVATE ${VULKAN_CXX_FLAGS})
 
-# vulkan_shader_lib
+# Find GLSL compiler executable
+
+if(ANDROID)
+  if(NOT ANDROID_NDK)
+    message(FATAL_ERROR "ANDROID_NDK not set")
+  endif()
+
+  set(GLSLC_PATH
+      "${ANDROID_NDK}/shader-tools/${ANDROID_NDK_HOST_SYSTEM_NAME}/glslc")
+else()
+  find_program(
+    GLSLC_PATH glslc
+    PATHS ENV VULKAN_SDK
+    PATHS "$ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/bin"
+    PATHS "$ENV{VULKAN_SDK}/bin")
+
+  if(NOT GLSLC_PATH)
+    message(FATAL_ERROR "USE_VULKAN glslc not found")
+  endif()
+endif()
+
+# Required to enable linking with --whole-archive
+include(${EXECUTORCH_ROOT}/build/Utils.cmake)
 
-file(GLOB VULKAN_IMPL_CPP ${ATEN_VULKAN_PATH}/impl/*.cpp)
+# Convenience macro to create a shader library
 
-add_library(vulkan_shader_lib STATIC ${VULKAN_IMPL_CPP} ${vulkan_generated_cpp})
+macro(vulkan_shader_library SHADERS_PATH LIBRARY_NAME)
+  set(VULKAN_SHADERGEN_ENV "")
+  set(VULKAN_SHADERGEN_OUT_PATH ${CMAKE_BINARY_DIR}/${LIBRARY_NAME})
 
-list(APPEND VULKAN_API_HEADERS ${CMAKE_BINARY_DIR}/vulkan)
+  execute_process(
+    COMMAND
+      "${PYTHON_EXECUTABLE}" ${PYTORCH_PATH}/tools/gen_vulkan_spv.py --glsl-path
+      ${SHADERS_PATH} --output-path ${VULKAN_SHADERGEN_OUT_PATH}
+      --glslc-path=${GLSLC_PATH} --tmp-dir-path=${VULKAN_SHADERGEN_OUT_PATH}
+      --env ${VULKAN_GEN_ARG_ENV}
+    RESULT_VARIABLE error_code)
+  set(ENV{PYTHONPATH} ${PYTHONPATH})
 
-target_include_directories(vulkan_shader_lib PRIVATE ${VULKAN_API_HEADERS})
+  set(vulkan_generated_cpp ${VULKAN_SHADERGEN_OUT_PATH}/spv.cpp)
 
-target_link_libraries(vulkan_shader_lib vulkan_api_lib)
+  add_library(${LIBRARY_NAME} STATIC ${vulkan_generated_cpp})
+  target_include_directories(${LIBRARY_NAME} PRIVATE ${COMMON_INCLUDES})
+  target_link_libraries(${LIBRARY_NAME} vulkan_api_lib)
+  target_compile_options(${LIBRARY_NAME} PRIVATE ${VULKAN_CXX_FLAGS})
+  # Link this library with --whole-archive due to dynamic shader registrations
+  target_link_options_shared_lib(${LIBRARY_NAME})
 
-target_compile_options(vulkan_shader_lib PRIVATE ${VULKAN_CXX_FLAGS})
+  unset(VULKAN_SHADERGEN_ENV)
+  unset(VULKAN_SHADERGEN_OUT_PATH)
+endmacro()

backends/vulkan/test/glsl/all_shaders.yaml

@@ -1,6 +1,13 @@
 binary_op_nobroadcast__test:
   parameter_names_with_default_values:
+    DTYPE: float
     OPERATOR: X + Y
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: "half"
+        SUFFIX: "half"
+      - VALUE: "float"
+        SUFFIX: "float"
   shader_variants:
     - NAME: binary_add_nobroadcast__test
       OPERATOR: X + Y
@@ -13,6 +20,19 @@ binary_op_nobroadcast__test:
     - NAME: binary_pow_nobroadcast__test
       OPERATOR: pow(X, Y)
 
+fill_texture__test:
+  parameter_names_with_default_values:
+    DTYPE: float
+    NDIM: 3
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: "half"
+        SUFFIX: "half"
+      - VALUE: "float"
+        SUFFIX: "float"
+  shader_variants:
+    - NAME: fill_texture__test
+
 image_to_nchw__test:
   parameter_names_with_default_values:
     NDIM: 3

backends/vulkan/test/glsl/binary_op_nobroadcast__test.glsl

@@ -9,15 +9,14 @@
 #version 450 core
 // clang-format off
 #define PRECISION ${PRECISION}
-#define FORMAT ${FORMAT}
 
 #define OP(X, Y) ${OPERATOR}
 // clang-format on
 
 layout(std430) buffer;
 
 // clang-format off
-layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D image_out;
+layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly image3D image_out;
 // clang-format on
 layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in;
 layout(set = 0, binding = 2) uniform PRECISION sampler3D image_other;

backends/vulkan/test/glsl/fill_texture__test.glsl

@@ -8,14 +8,13 @@
 
 #version 450 core
 #define PRECISION ${PRECISION}
-#define FORMAT ${FORMAT}
 
 layout(std430) buffer;
 
 /* Qualifiers: layout - storage - precision - memory */
 
 // clang-format off
-layout(set = 0, binding = 0, FORMAT) uniform PRECISION restrict writeonly image3D uOutput;
+layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} uOutput;
 // clang-format on
 layout(set = 0, binding = 1) uniform PRECISION restrict Block {
   ivec3 size;

backends/vulkan/test/utils/test_utils.cpp

@@ -115,15 +115,19 @@ void record_image_to_nchw_op(
       v_src.cpu_sizes_ubo()->buffer());
 }
 
-void record_arithmetic_op(
+void record_binary_op(
     api::Context* const context,
-    const api::ShaderInfo& compute_shader,
+    const std::string& op_name,
     vTensor& v_in1,
     vTensor& v_in2,
     vTensor& v_dst) {
+  std::stringstream kernel_name;
+  kernel_name << "binary_" << op_name << "_nobroadcast__test";
+  apply_dtype_suffix(kernel_name, v_dst);
+
   api::PipelineBarrier pipeline_barrier{};
   context->submit_compute_job(
-      compute_shader,
+      VK_KERNEL_FROM_STR(kernel_name.str()),
       pipeline_barrier,
       v_dst.virtual_extents(),
       adaptive_work_group_size(v_dst.virtual_extents()),
@@ -144,14 +148,17 @@ void execute_and_check_add(
     float a_val,
     float b_val) {
   // Add shader kernel
-  api::ShaderInfo kernel = VK_KERNEL(binary_add_nobroadcast__test);
+  api::ShaderInfo kernel = VK_KERNEL(binary_add_nobroadcast__test_half);
+  if (c.image().format() == VK_FORMAT_R32G32B32A32_SFLOAT) {
+    kernel = VK_KERNEL(nchw_to_image3d__test_C_packed_float);
+  }
 
   // Fill input tensors
   fill_vtensor(a, a_val);
   fill_vtensor(b, b_val);
 
   // a + b = c
-  record_arithmetic_op(api::context(), kernel, a, b, c);
+  record_binary_op(api::context(), "add", a, b, c);
 
   // Extract output tensor
   std::vector<float> data_out = extract_vtensor(c);

backends/vulkan/test/utils/test_utils.h

@@ -12,6 +12,7 @@
 
 #include <ATen/native/vulkan/api/api.h>
 
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/StagingUtils.h>
 
 using namespace at::native::vulkan;
@@ -70,9 +71,9 @@ void record_image_to_nchw_op(
     vTensor& v_src,
     api::VulkanBuffer& dst_buffer);
 
-void record_arithmetic_op(
+void record_binary_op(
     api::Context* const context,
-    const api::ShaderInfo& compute_shader,
+    const std::string& op_name,
     vTensor& v_in1,
     vTensor& v_in2,
     vTensor& v_dst);