[ET-VK] Enable auto-generated operator correctness tests and benchmark binaries in OSS (#10260)

Summary:
## Context

As title. The majority of operator correctness testing for the Vulkan
delegate is done via the tests under the `op_tests/` directory, which
includes generated operator correctness and benchmark tests. However,
these tests are currently only able to be built and run via buck, and
are only tested in the Meta internal repo.

This PR sets up the CMake build for the operator tests under
`op_tests/`, allowing them to be built and run in the Github repo.

## Next Steps

* Run these tests in CI

This work is in service of making it easier for Open Source
Constributors to contribute to the ExecuTorch Vulkan Delegate.

Test Plan:
## Test Plan

Build and run the test binaries on a Linux machine.

### Setup

```
# Set up ExecuTorch
cd ~/executorch
./install_executorch.sh

# Build + Install C++ components
rm -rf cmake-out && \
  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_DEVTOOLS=ON \
  -DEXECUTORCH_BUILD_VULKAN=ON \
  -DEXECUTORCH_BUILD_XNNPACK=ON \
  -DEXECUTORCH_BUILD_TESTS=ON \
  -Bcmake-out && \
  cmake --build cmake-out -j64 --target install
```

### Build tests

```
rm -rf cmake-out/backends/vulkan/test/op_tests && \
  pp cmake backends/vulkan/test/op_tests \
    -DCMAKE_INSTALL_PREFIX=cmake-out \
    -DPYTHON_EXECUTABLE=python \
    -DTORCH_OPS_YAML_PATH=/home/ssjia/Github/pytorch/aten/src/ATen/native \
    -Bcmake-out/backends/vulkan/test/op_tests && \
  cmake --build cmake-out/backends/vulkan/test/op_tests -j16
```

Note: the auto-generated operator correctness tests need the `tags.yaml`
and `native_functions.yaml` file from the pytorch repo. Otherwise they
won't be built.

### Run tests

```
cmake-out/backends/vulkan/test/op_tests/vulkan_sdpa_test
cmake-out/backends/vulkan/test/op_tests/vulkan_rope_test
cmake-out/backends/vulkan/test/op_tests/vulkan_linear_weight_int4_test
cmake-out/backends/vulkan/test/op_tests/vulkan_operator_correctness_tests --gtest_filter="*add_Tensor*"
```

Author: SS-JIA

backends/vulkan/cmake/ShaderLibrary.cmake

@@ -45,16 +45,20 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
 
 function(gen_vulkan_shader_lib_cpp shaders_path)
   set(VULKAN_SHADERGEN_ENV "")
-  set(VULKAN_SHADERGEN_OUT_PATH ${CMAKE_BINARY_DIR}/${ARGV1})
+  set(VULKAN_SHADERGEN_OUT_PATH ${CMAKE_BINARY_DIR}/vulkan_compute_shaders)
 
-  execute_process(
+  add_custom_command(
+    COMMENT "Generating Vulkan Compute Shaders"
+    OUTPUT ${VULKAN_SHADERGEN_OUT_PATH}/spv.cpp
     COMMAND
       "${PYTHON_EXECUTABLE}"
       ${EXECUTORCH_ROOT}/backends/vulkan/runtime/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}/shader_cache/
-      --env ${VULKAN_GEN_ARG_ENV}
-    RESULT_VARIABLE error_code
+      --glslc-path=${GLSLC_PATH}
+      --tmp-dir-path=${VULKAN_SHADERGEN_OUT_PATH}/shader_cache/ --env
+      ${VULKAN_GEN_ARG_ENV}
+    DEPENDS ${shaders_path}/*
+            ${EXECUTORCH_ROOT}/backends/vulkan/runtime/gen_vulkan_spv.py
   )
 
   set(generated_spv_cpp
@@ -86,13 +90,6 @@ macro(vulkan_shader_library shaders_path library_name)
   set(VULKAN_SHADERGEN_ENV "")
   set(VULKAN_SHADERGEN_OUT_PATH ${CMAKE_BINARY_DIR}/${library_name})
 
-  # execute_process( COMMAND "${PYTHON_EXECUTABLE}"
-  # ${EXECUTORCH_ROOT}/backends/vulkan/runtime/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})
-
   set(generated_spv_cpp ${VULKAN_SHADERGEN_OUT_PATH}/spv.cpp)
 
   add_library(${library_name} STATIC ${generated_spv_cpp})

backends/vulkan/test/CMakeLists.txt

@@ -46,8 +46,7 @@ if(LIB_VULKAN_BACKEND)
   set(VULKAN_THIRD_PARTY_PATH ${CMAKE_CURRENT_SOURCE_DIR}/../third-party)
 
   set(GTEST_INCLUDE_PATH
-      ${EXECUTORCH_ROOT}/third-party/googletest/googletest/include set
-      (PYTORCH_PATH ${EXECUTORCH_ROOT}/third-party/pytorch)
+      ${EXECUTORCH_ROOT}/third-party/googletest/googletest/include
   )
   set(VULKAN_HEADERS_PATH ${VULKAN_THIRD_PARTY_PATH}/Vulkan-Headers/include)
   set(VOLK_PATH ${VULKAN_THIRD_PARTY_PATH}/volk)

backends/vulkan/test/op_tests/CMakeLists.txt

@@ -0,0 +1,166 @@
+# 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.
+
+# ### Editing this file ###
+#
+# This file should be formatted with
+# ~~~
+# cmake-format -i CMakeLists.txt
+# ~~~
+# It should also be cmake-lint clean.
+#
+# The targets in this file will be built if EXECUTORCH_BUILD_VULKAN is ON
+
+cmake_minimum_required(VERSION 3.19)
+project(executorch)
+
+find_package(executorch CONFIG REQUIRED COMPONENTS vulkan_backend)
+find_package(GTest CONFIG REQUIRED)
+
+if(NOT EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..)
+endif()
+
+# 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}/tools/cmake/Utils.cmake)
+
+get_torch_base_path(TORCH_BASE_PATH)
+message(STATUS "torch base path: ${TORCH_BASE_PATH}")
+
+# Only build tests if Vulkan was compiled
+find_library(LIB_VULKAN_BACKEND vulkan_backend)
+find_library(LIB_TORCH torch ${TORCH_BASE_PATH}/lib)
+find_library(LIB_TORCH_CPU torch_cpu ${TORCH_BASE_PATH}/lib)
+find_library(LIB_C10 c10 ${TORCH_BASE_PATH}/lib)
+
+message(STATUS "Vulkan backend lib ${LIB_VULKAN_BACKEND}")
+message(STATUS "Torch ${LIB_TORCH}")
+
+if(NOT PYTHON_EXECUTABLE)
+  set(PYTHON_EXECUTABLE python3)
+endif()
+
+# Third party include paths
+
+set(VULKAN_THIRD_PARTY_PATH ${CMAKE_CURRENT_SOURCE_DIR}/../../third-party)
+
+set(GTEST_INCLUDE_PATH
+    ${EXECUTORCH_ROOT}/third-party/googletest/googletest/include
+)
+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)
+
+set(COMMON_INCLUDES
+    ${EXECUTORCH_ROOT}/..
+    ${VULKAN_HEADERS_PATH}
+    ${VOLK_PATH}
+    ${VMA_PATH}
+    ${GTEST_INCLUDE_PATH}
+    ${TORCH_BASE_PATH}/include
+    ${TORCH_BASE_PATH}/include/torch/csrc/api/include
+)
+
+target_link_options_shared_lib(vulkan_backend)
+
+function(vulkan_op_test test_name test_src)
+  set(extra_deps ${ARGN})
+
+  add_executable(${test_name} ${test_src})
+  target_include_directories(${test_name} PRIVATE ${COMMON_INCLUDES})
+  target_link_libraries(
+    ${test_name}
+    PRIVATE GTest::gtest_main
+            vulkan_backend
+            executorch
+            ${LIB_TORCH}
+            ${LIB_TORCH_CPU}
+            ${LIB_C10}
+            ${extra_deps}
+  )
+
+  add_test(${test_name} ${test_name})
+endfunction()
+
+if(LIB_VULKAN_BACKEND AND LIB_TORCH)
+  find_library(
+    CUSTOM_OPS_LIB custom_ops_aot_lib
+    HINTS ${CMAKE_INSTALL_PREFIX}/executorch/extension/llm/custom_ops
+  )
+  if(CUSTOM_OPS_LIB)
+    vulkan_op_test(
+      vulkan_sdpa_test ${CMAKE_CURRENT_SOURCE_DIR}/sdpa_test.cpp
+      ${CUSTOM_OPS_LIB}
+    )
+  else()
+    message(
+      STATUS "Skip building sdpa_test because custom_ops_aot_lib is not found"
+    )
+  endif()
+  vulkan_op_test(
+    vulkan_rope_test ${CMAKE_CURRENT_SOURCE_DIR}/rotary_embedding_test.cpp
+  )
+  vulkan_op_test(
+    vulkan_linear_weight_int4_test
+    ${CMAKE_CURRENT_SOURCE_DIR}/linear_weight_int4_test.cpp
+  )
+
+  # Only build generated op tests if a path to tags.yaml and
+  # native_functions.yaml is provided. These files are required for codegen.
+  if(TORCH_OPS_YAML_PATH)
+    set(GENERATED_VULKAN_TESTS_CPP_PATH ${CMAKE_CURRENT_BINARY_DIR}/vk_gen_cpp)
+
+    # Generated operator correctness tests
+
+    set(generated_test_cpp ${GENERATED_VULKAN_TESTS_CPP_PATH}/op_tests.cpp)
+
+    add_custom_command(
+      COMMENT "Generating Vulkan operator correctness tests"
+      OUTPUT ${generated_test_cpp}
+      COMMAND
+        ${PYTHON_EXECUTABLE}
+        ${EXECUTORCH_ROOT}/backends/vulkan/test/op_tests/generate_op_correctness_tests.py
+        -o ${GENERATED_VULKAN_TESTS_CPP_PATH} --tags-path
+        ${TORCH_OPS_YAML_PATH}/tags.yaml --aten-yaml-path
+        ${TORCH_OPS_YAML_PATH}/native_functions.yaml
+      DEPENDS ${EXECUTORCH_ROOT}/backends/vulkan/test/op_tests/**/*.py
+    )
+
+    vulkan_op_test(vulkan_op_correctness_tests ${generated_test_cpp})
+
+    # Generated operator benchmarks (only built in google benchmark is
+    # installed)
+    find_package(benchmark CONFIG)
+
+    if(benchmark_FOUND)
+      set(generated_benchmark_cpp
+          ${GENERATED_VULKAN_TESTS_CPP_PATH}/op_benchmarks.cpp
+      )
+
+      add_custom_command(
+        COMMENT "Generating Vulkan operator benchmarks"
+        OUTPUT ${generated_benchmark_cpp}
+        COMMAND
+          ${PYTHON_EXECUTABLE}
+          ${EXECUTORCH_ROOT}/backends/vulkan/test/op_tests/generate_op_benchmarks.py
+          -o ${GENERATED_VULKAN_TESTS_CPP_PATH} --tags-path
+          ${TORCH_OPS_YAML_PATH}/tags.yaml --aten-yaml-path
+          ${TORCH_OPS_YAML_PATH}/native_functions.yaml
+        DEPENDS ${EXECUTORCH_ROOT}/backends/vulkan/test/op_tests/**/*.py
+      )
+
+      vulkan_op_test(vulkan_op_benchmarks ${generated_benchmark_cpp})
+    endif()
+  else()
+    message(
+      STATUS
+        "Skipping generated operator correctness tests and benchmarks. Please specify TORCH_OPS_YAML_PATH to build these tests."
+    )
+  endif()
+endif()

backends/vulkan/test/op_tests/generate_op_correctness_tests.py

@@ -58,6 +58,9 @@ def process_test_suites(
 def generate_cpp(
     native_functions_yaml_path: str, tags_path: str, output_dir: str
 ) -> None:
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
     output_file = os.path.join(output_dir, "op_tests.cpp")
     cpp_generator = VkCorrectnessTestFileGen(output_file)
 

backends/vulkan/test/op_tests/utils/gen_benchmark_vk.py

@@ -228,15 +228,15 @@ def generate_benchmark_fixture(self) -> str:
   return at::from_blob(values.data(), sizes, at::kFloat).toType(dtype).detach().clone();
 }}
 
-at::Tensor make_index_tensor(std::vector<int64_t> indices) {{
+at::Tensor make_index_tensor_1d(std::vector<int64_t> indices) {{
   at::ScalarType dtype = at::kInt;
   std::vector<int64_t> sizes = {{static_cast<int64_t>(indices.size())}};
 
   // Clone as original data will be deallocated upon return.
   return at::from_blob(indices.data(), sizes, dtype).detach().clone();
 }}
 
-at::Tensor make_index_tensor(std::vector<std::vector<int64_t>> indices) {{
+at::Tensor make_index_tensor_2d(std::vector<std::vector<int64_t>> indices) {{
   at::ScalarType dtype = at::kInt;
   std::vector<int64_t> sizes = {{
     static_cast<int64_t>(indices.size()),
@@ -252,7 +252,7 @@ def generate_benchmark_fixture(self) -> str:
   return at::from_blob(acc.data(), sizes, dtype).detach().clone();
 }}
 
-at::Tensor make_index_tensor(std::vector<std::vector<std::vector<int64_t>>> indices) {{
+at::Tensor make_index_tensor_3d(std::vector<std::vector<std::vector<int64_t>>> indices) {{
   at::ScalarType dtype = at::kInt;
   std::vector<int64_t> sizes = {{
     static_cast<int64_t>(indices.size()),

backends/vulkan/test/op_tests/utils/gen_computegraph.py

@@ -229,11 +229,10 @@ def create_aten_fn_call(self) -> str:
 
     def create_aten_method_call(self) -> str:
         # For functions with only Method variant, we fallback to the function
-        # declared in MethodOperators.h. The method is declared as
-        # at::_ops::{name}::call(*), and ATEN_FN is a handly macro.
+        # declared in MethodOperators.h
         cpp_sig = gen_static_dispatch_backend_call_signature(self.f_sig, self.f)
         exprs = translate_args(self.f_sig, cpp_sig)
-        func_call = f"ATEN_FN({self.f_sig.name()})({exprs});"
+        func_call = f"at::_ops::{self.f_sig.name()}::call({exprs});"
         return func_call
 
     def create_out_src(self, include_declarations: bool = True) -> str:

backends/vulkan/test/op_tests/utils/gen_correctness_base.py

@@ -170,7 +170,13 @@ def create_input_data(self, arg: Argument, data: Any) -> str:  # noqa: C901
 
         if cpp_type == AT_TENSOR:
             if arg.name == "index" or arg.name == "indices":
-                ret_str += f"make_index_tensor({init_list_str(data)});"
+                args_str = init_list_str(data)
+                if args_str[:3] == "{{{":
+                    ret_str += f"make_index_tensor_3d({init_list_str(data)});"
+                elif args_str[:2] == "{{":
+                    ret_str += f"make_index_tensor_2d({init_list_str(data)});"
+                else:
+                    ret_str += f"make_index_tensor_1d({init_list_str(data)});"
             else:
                 ret_str += self.call_data_gen_fn(arg, data)
         elif cpp_type == OPT_AT_TENSOR:
@@ -278,7 +284,7 @@ def generate_suite_cpp(self) -> str:
     float high = 1.0) {{
   if (high == 1.0 && low == 0.0)
     return at::rand(sizes, at::device(at::kCPU).dtype(dtype));
-    
+
   if (dtype == at::kChar)
     return at::randint(high, sizes, at::device(at::kCPU).dtype(dtype));
 
@@ -307,15 +313,15 @@ def generate_suite_cpp(self) -> str:
   return at::from_blob(values.data(), sizes, at::kFloat).toType(dtype).detach().clone();
 }}
 
-at::Tensor make_index_tensor(std::vector<int64_t> indices) {{
+at::Tensor make_index_tensor_1d(std::vector<int64_t> indices) {{
   at::ScalarType dtype = at::kInt;
   std::vector<int64_t> sizes = {{static_cast<int64_t>(indices.size())}};
 
   // Clone as original data will be deallocated upon return.
   return at::from_blob(indices.data(), sizes, dtype).detach().clone();
 }}
 
-at::Tensor make_index_tensor(std::vector<std::vector<int64_t>> indices) {{
+at::Tensor make_index_tensor_2d(std::vector<std::vector<int64_t>> indices) {{
   at::ScalarType dtype = at::kInt;
   std::vector<int64_t> sizes = {{
     static_cast<int64_t>(indices.size()),
@@ -331,7 +337,7 @@ def generate_suite_cpp(self) -> str:
   return at::from_blob(acc.data(), sizes, dtype).detach().clone();
 }}
 
-at::Tensor make_index_tensor(std::vector<std::vector<std::vector<int64_t>>> indices) {{
+at::Tensor make_index_tensor_3d(std::vector<std::vector<std::vector<int64_t>>> indices) {{
   at::ScalarType dtype = at::kInt;
   std::vector<int64_t> sizes = {{
     static_cast<int64_t>(indices.size()),