[ET-VK] Fix OSS build for Vulkan Delegate

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);