Add Vulkan backend to llama demo binary (#2608)

Summary:
Pull Request resolved: #2608

## Context

Cleans up Vulkan backend CMake build so that `vulkan_backend` target is a single static library will all required symbols. Then add vulkan backend to llama runner binary.

imported-using-ghimport

Test Plan: Imported from OSS

Reviewed By: mcr229, cbilgin

Differential Revision: D55272794

Pulled By: SS-JIA

fbshipit-source-id: 16074d32315a01df0a2fd70311ae98c07da12dc9

Author: SS-JIA

backends/vulkan/CMakeLists.txt

@@ -49,21 +49,23 @@ set(COMMON_INCLUDES ${VULKAN_API_HEADERS} ${EXECUTORCH_ROOT}/..)
 # vulkan_graph_lib
 
 file(GLOB_RECURSE vulkan_graph_cpp ${RUNTIME_PATH}/graph/*)
+list(APPEND vulkan_graph_cpp ${vulkan_api_cpp})
 
 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)
 
-# Due to dynamic registrations, these libraries must be explicitly linked
-set(VULKAN_STANDARD_OPS_LIBS vulkan_graph_lib vulkan_graph_shaderlib)
-
-# vulkan_graph_shaderlib
+# vulkan_standard_shaders
 
 set(VULKAN_GRAPH_SHADERS_PATH ${RUNTIME_PATH}/graph/ops/glsl/)
-vulkan_shader_library(${VULKAN_GRAPH_SHADERS_PATH} vulkan_graph_shaderlib)
+# Declares a vulkan_standard_shaders library containing compiled SPIR-V shaders
+# from the above path
+vulkan_shader_library(${VULKAN_GRAPH_SHADERS_PATH} vulkan_standard_shaders)
+# preserve the path of the generated cpp file
+set(vulkan_standard_shaders_cpp
+    ${CMAKE_BINARY_DIR}/vulkan_standard_shaders/spv.cpp)
 
 # Generate Files from flatc
 
@@ -90,15 +92,16 @@ target_include_directories(
   vulkan_schema INTERFACE ${SCHEMA_INCLUDE_DIR}
                           ${EXECUTORCH_ROOT}/third-party/flatbuffers/include)
 
-# vulkan_backend_lib
+# vulkan_backend
 
 file(GLOB vulkan_backend_cpp ${RUNTIME_PATH}/*.cpp)
+list(APPEND vulkan_backend_cpp ${vulkan_graph_cpp})
+list(APPEND vulkan_backend_cpp ${vulkan_standard_shaders_cpp})
 
-add_library(vulkan_backend ${vulkan_backend_cpp})
+add_library(vulkan_backend STATIC ${vulkan_backend_cpp})
 target_include_directories(vulkan_backend PRIVATE ${SCHEMA_INCLUDE_DIR}
                                                   ${COMMON_INCLUDES})
-target_link_libraries(vulkan_backend PRIVATE vulkan_graph_lib vulkan_schema
-                                             executorch)
+target_link_libraries(vulkan_backend PRIVATE vulkan_schema executorch)
 target_compile_options(vulkan_backend PRIVATE ${VULKAN_CXX_FLAGS})
 # Link this library with --whole-archive due to dynamic backend registration
 target_link_options_shared_lib(vulkan_backend)
@@ -110,15 +113,13 @@ if(NOT CMAKE_TOOLCHAIN_FILE MATCHES ".*iOS\.cmake$")
   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} vulkan_schema
-    vulkan_backend ${VULKAN_STANDARD_OPS_LIBS})
+  target_link_libraries(vulkan_executor_runner ${_executor_runner_libs}
+                        vulkan_schema vulkan_backend)
   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} vulkan_schema
-    vulkan_backend vulkan_api_lib ${VULKAN_STANDARD_OPS_LIBS})
+  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()
 
@@ -140,7 +141,13 @@ if(EXECUTORCH_BUILD_GTESTS)
   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)
+    vulkan_compute_api_test PRIVATE gtest_main vulkan_graph_lib
+                                    vulkan_standard_shaders test_shaderlib)
   target_compile_options(vulkan_compute_api_test PRIVATE ${VULKAN_CXX_FLAGS})
 endif()
+
+install(
+  TARGETS vulkan_backend
+  DESTINATION lib
+  INCLUDES
+  DESTINATION ${COMMON_INCLUDES})

backends/vulkan/cmake/ATenVulkan.cmake

@@ -22,7 +22,7 @@ if(NOT VULKAN_THIRD_PARTY_PATH)
   set(VULKAN_THIRD_PARTY_PATH ${CMAKE_CURRENT_SOURCE_DIR}/../third-party)
 endif()
 
-# Source paths and compile settings
+# Source file paths
 
 set(ATEN_PATH ${PYTORCH_PATH}/aten/src)
 set(ATEN_VULKAN_PATH ${ATEN_PATH}/ATen/native/vulkan)
@@ -31,28 +31,26 @@ set(VULKAN_HEADERS_PATH ${VULKAN_THIRD_PARTY_PATH}/Vulkan-Headers/include)
 set(VOLK_PATH ${VULKAN_THIRD_PARTY_PATH}/volk)
 set(VMA_PATH ${VULKAN_THIRD_PARTY_PATH}/VulkanMemoryAllocator)
 
+# Compile settings
+
 set(VULKAN_CXX_FLAGS "")
 list(APPEND VULKAN_CXX_FLAGS "-DUSE_VULKAN_API")
 list(APPEND VULKAN_CXX_FLAGS "-DUSE_VULKAN_WRAPPER")
 list(APPEND VULKAN_CXX_FLAGS "-DUSE_VULKAN_VOLK")
 list(APPEND VULKAN_CXX_FLAGS "-DVK_NO_PROTOTYPES")
 list(APPEND VULKAN_CXX_FLAGS "-DVOLK_DEFAULT_VISIBILITY")
 
-# vulkan_api_lib
-
-file(GLOB VULKAN_API_CPP ${ATEN_VULKAN_PATH}/api/*.cpp)
+# vulkan_api source files
 
-add_library(vulkan_api_lib STATIC ${VULKAN_API_CPP} ${VOLK_PATH}/volk.c)
+file(GLOB vulkan_api_cpp ${ATEN_VULKAN_PATH}/api/*.cpp)
+list(APPEND vulkan_api_cpp ${VOLK_PATH}/volk.c)
 
 set(VULKAN_API_HEADERS)
 list(APPEND VULKAN_API_HEADERS ${ATEN_PATH})
 list(APPEND VULKAN_API_HEADERS ${VULKAN_HEADERS_PATH})
 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})
-
 # Find GLSL compiler executable
 
 if(ANDROID)
@@ -65,9 +63,8 @@ if(ANDROID)
 else()
   find_program(
     GLSLC_PATH glslc
-    PATHS ENV VULKAN_SDK
-    PATHS "$ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/bin"
-    PATHS "$ENV{VULKAN_SDK}/bin")
+    PATHS ENV VULKAN_SDK "$ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/bin"
+          "$ENV{VULKAN_SDK}/bin")
 
   if(NOT GLSLC_PATH)
     message(FATAL_ERROR "USE_VULKAN glslc not found")
@@ -77,29 +74,31 @@ endif()
 # Required to enable linking with --whole-archive
 include(${EXECUTORCH_ROOT}/build/Utils.cmake)
 
-# Convenience macro to create a shader library
-
-macro(vulkan_shader_library SHADERS_PATH LIBRARY_NAME)
+# Convenience macro to generate a SPIR-V shader library target. Given the path
+# to the shaders to compile and the name of the library, it will create a static
+# library containing the generated SPIR-V shaders. The generated_spv_cpp
+# variable can be used to reference the generated CPP file outside the macro.
+macro(VULKAN_SHADER_LIBRARY shaders_path library_name)
   set(VULKAN_SHADERGEN_ENV "")
-  set(VULKAN_SHADERGEN_OUT_PATH ${CMAKE_BINARY_DIR}/${LIBRARY_NAME})
+  set(VULKAN_SHADERGEN_OUT_PATH ${CMAKE_BINARY_DIR}/${library_name})
 
   execute_process(
     COMMAND
       "${PYTHON_EXECUTABLE}" ${PYTORCH_PATH}/tools/gen_vulkan_spv.py --glsl-path
-      ${SHADERS_PATH} --output-path ${VULKAN_SHADERGEN_OUT_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})
 
-  set(vulkan_generated_cpp ${VULKAN_SHADERGEN_OUT_PATH}/spv.cpp)
+  set(generated_spv_cpp ${VULKAN_SHADERGEN_OUT_PATH}/spv.cpp)
 
-  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})
+  add_library(${library_name} STATIC ${generated_spv_cpp})
+  target_include_directories(${library_name} PRIVATE ${COMMON_INCLUDES})
+  target_link_libraries(${library_name} vulkan_graph_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_link_options_shared_lib(${library_name})
 
   unset(VULKAN_SHADERGEN_ENV)
   unset(VULKAN_SHADERGEN_OUT_PATH)

build/executorch-config.cmake

@@ -44,7 +44,7 @@ endif()
 set(lib_list
     etdump bundled_program extension_data_loader ${FLATCC_LIB} mpsdelegate
     qnn_executorch_backend portable_ops_lib extension_module xnnpack_backend
-    XNNPACK cpuinfo pthreadpool
+    XNNPACK cpuinfo pthreadpool vulkan_backend
 )
 foreach(lib ${lib_list})
     # Name of the variable which stores result of the find_library search

examples/models/llama2/CMakeLists.txt

@@ -69,6 +69,16 @@ if(TARGET xnnpack_backend)
                       $<TARGET_FILE:xnnpack_backend> \
                       LINKER:--no-whole-archive")
 endif()
+
+# Vulkan backend
+if(TARGET vulkan_backend)
+  target_link_libraries(llama_main PUBLIC vulkan_backend)
+  target_link_options(
+    llama_main PUBLIC "SHELL:LINKER:--whole-archive \
+                      $<TARGET_FILE:vulkan_backend> \
+                      LINKER:--no-whole-archive")
+endif()
+
 target_compile_options(llama_main PUBLIC ${_common_compile_options})
 
 # Print all summary

examples/models/llama2/export_llama_lib.py

@@ -577,7 +577,7 @@ def _export_llama(modelname, args) -> str:  # noqa: C901
 
     if args.vulkan:
         assert (
-            args.dtype_override is None
+            args.dtype_override == "fp32" or args.dtype_override is None
         ), "Vulkan backend does not support non fp32 dtypes at the moment"
         assert (
             args.quantization_mode is None