Cmake sdpa custom op (#2671)

kimishpatel · facebook-github-bot · commit 3881bff40a91 · 2024-03-27T21:44:26.000-07:00
Summary: Pull Request resolved: #2671 Enable `--use_sdpa_with_kv_cache` with Cmake. ghstack-source-id: 220378207 exported-using-ghexport Pre-existing lint failures bypass-github-export-checks bypass-github-pytorch-ci-checks Reviewed By: digantdesai, mcr229 Differential Revision: D55344016 fbshipit-source-id: 7dafab86dc66da85c7a988a728bb5cce5382e3d6
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -45,6 +45,26 @@ endif()
 if(NOT CMAKE_BUILD_TYPE)
   set(CMAKE_BUILD_TYPE Debug)
 endif()
+
+# --- cpuinfo
+set(CPUINFO_SOURCE_DIR "backends/xnnpack/third-party/cpuinfo")
+set(CPUINFO_BUILD_TOOLS OFF CACHE BOOL "")
+set(CPUINFO_BUILD_UNIT_TESTS OFF CACHE BOOL "")
+set(CPUINFO_BUILD_MOCK_TESTS OFF CACHE BOOL "")
+set(CPUINFO_BUILD_BENCHMARKS OFF CACHE BOOL "")
+set(CPUINFO_LIBRARY_TYPE "static" CACHE STRING "")
+set(CPUINFO_LOG_LEVEL "error" CACHE STRING "")
+set(CLOG_SOURCE_DIR "${CPUINFO_SOURCE_DIR}/deps/clog")
+add_subdirectory("${CPUINFO_SOURCE_DIR}")
+
+# --- pthreadpool
+set(PTHREADPOOL_SOURCE_DIR "backends/xnnpack/third-party/pthreadpool")
+set(PTHREADPOOL_BUILD_TESTS OFF CACHE BOOL "")
+set(PTHREADPOOL_BUILD_BENCHMARKS OFF CACHE BOOL "")
+set(PTHREADPOOL_LIBRARY_TYPE "static" CACHE STRING "")
+set(PTHREADPOOL_ALLOW_DEPRECATED_API ON CACHE BOOL "")
+add_subdirectory("${PTHREADPOOL_SOURCE_DIR}")
+
 # ------------------------------ OPTIONS -------------------------------------
 # WARNING: Please don't add example specific options in this CMakeLists.txt.
 # Instead please use `find_package(executorch REQUIRED)` in the example
diff --git a/backends/xnnpack/CMakeLists.txt b/backends/xnnpack/CMakeLists.txt
@@ -87,6 +87,10 @@ target_link_libraries(xnnpack_backend
 target_include_directories(xnnpack_backend
                            PUBLIC ${_common_include_directories})
 target_include_directories(xnnpack_backend PUBLIC ${XNNPACK_INCLUDE_DIR})
+target_include_directories(xnnpack_backend PUBLIC
+${CMAKE_CURRENT_SOURCE_DIR}/third-party/pthreadpool/include)
+target_include_directories(xnnpack_backend PUBLIC
+${CMAKE_CURRENT_SOURCE_DIR}/third-party/cpuinfo/include)
 target_compile_options(xnnpack_backend PUBLIC ${_common_compile_options})
 target_link_options_shared_lib(xnnpack_backend)
 
@@ -101,7 +105,8 @@ if(NOT CMAKE_TOOLCHAIN_FILE MATCHES ".*iOS\.cmake$")
   list(TRANSFORM _xnn_executor_runner__srcs PREPEND "${EXECUTORCH_ROOT}/")
   add_executable(xnn_executor_runner ${_xnn_executor_runner__srcs})
   target_link_libraries(xnn_executor_runner
-                        xnnpack_backend gflags portable_ops_lib)
+                        xnnpack_backend gflags portable_ops_lib
+                        pthreadpool cpuinfo)
   target_compile_options(xnn_executor_runner PUBLIC ${_common_compile_options})
 endif()
 
diff --git a/backends/xnnpack/cmake/Dependencies.cmake b/backends/xnnpack/cmake/Dependencies.cmake
@@ -8,27 +8,6 @@
 # https://github.com/pytorch/pytorch/blob/main/cmake/Dependencies.cmake
 set(THIRD_PARTY_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/third-party")
 
-# --- cpuinfo
-set(CPUINFO_SOURCE_DIR "${THIRD_PARTY_ROOT}/cpuinfo")
-set(CPUINFO_BUILD_TOOLS OFF CACHE BOOL "")
-set(CPUINFO_BUILD_UNIT_TESTS OFF CACHE BOOL "")
-set(CPUINFO_BUILD_MOCK_TESTS OFF CACHE BOOL "")
-set(CPUINFO_BUILD_BENCHMARKS OFF CACHE BOOL "")
-set(CPUINFO_LIBRARY_TYPE "static" CACHE STRING "")
-set(CPUINFO_LOG_LEVEL "error" CACHE STRING "")
-set(CLOG_SOURCE_DIR "${CPUINFO_SOURCE_DIR}/deps/clog")
-add_subdirectory("${CPUINFO_SOURCE_DIR}")
-list(APPEND xnnpack_third_party cpuinfo)
-
-# --- pthreadpool
-set(PTHREADPOOL_SOURCE_DIR "${THIRD_PARTY_ROOT}/pthreadpool")
-set(PTHREADPOOL_BUILD_TESTS OFF CACHE BOOL "")
-set(PTHREADPOOL_BUILD_BENCHMARKS OFF CACHE BOOL "")
-set(PTHREADPOOL_LIBRARY_TYPE "static" CACHE STRING "")
-set(PTHREADPOOL_ALLOW_DEPRECATED_API ON CACHE BOOL "")
-add_subdirectory("${PTHREADPOOL_SOURCE_DIR}")
-list(APPEND xnnpack_third_party pthreadpool)
-
 # --- XNNPACK
 set(XNNPACK_SOURCE_DIR "${THIRD_PARTY_ROOT}/XNNPACK")
 set(XNNPACK_INCLUDE_DIR "${XNNPACK_SOURCE_DIR}/include")
diff --git a/build/cmake_deps.toml b/build/cmake_deps.toml
@@ -265,6 +265,21 @@ filters = [
 ]
 # ---------------------------------- Vulkan end -----------------------------------
 # ---------------------------------- LLama start ----------------------------------
+[targets.custom_ops]
+buck_targets = [
+  "//examples/models/llama2/custom_ops:custom_ops",
+]
+filters = [
+  ".cpp$",
+]
+excludes = [
+  "^codegen",
+  # Exclude blas, since it's built as a separate target.
+  # "^kernels/optimized/blas",
+]
+deps = [
+  "executorch",
+]
 
 [targets.llama_runner]
 buck_targets = [
@@ -277,6 +292,7 @@ excludes = [
   "^codegen",
 ]
 deps = [
+  "custom_ops",
   "executorch",
   "extension_data_loader",
   "extension_module",
diff --git a/examples/models/llama2/CMakeLists.txt b/examples/models/llama2/CMakeLists.txt
@@ -45,17 +45,26 @@ find_package(gflags REQUIRED PATHS
 #
 # llama_main: test binary to run llama, with tokenizer and sampler integrated
 #
-add_executable(llama_main main.cpp)
+add_executable(llama_main main.cpp
+${CMAKE_CURRENT_SOURCE_DIR}/../../../backends/xnnpack/threadpool/cpuinfo_utils.cpp)
 if(CMAKE_BUILD_TYPE EQUAL "RELEASE")
   target_link_options(llama_main PRIVATE "LINKER:--gc-sections")
 endif()
 
 # find `executorch` libraries
 find_package(executorch CONFIG REQUIRED)
 
+# custom ops library
+add_subdirectory(custom_ops)
+
 # llama_runner library
 add_subdirectory(runner)
 
+target_include_directories(llama_main PUBLIC
+${CMAKE_CURRENT_SOURCE_DIR}/../../../backends/xnnpack/third-party/cpuinfo/include)
+target_include_directories(llama_main PUBLIC
+${CMAKE_CURRENT_SOURCE_DIR}/../../../backends/xnnpack/third-party/pthreadpool/include)
+
 set(link_libraries)
 
 if(EXECUTORCH_BUILD_OPTIMIZED)
@@ -67,7 +76,7 @@ else()
   target_link_options_shared_lib(portable_ops_lib)
 endif()
 
-target_link_libraries(llama_main PUBLIC gflags llama_runner)
+target_link_libraries(llama_main PUBLIC gflags llama_runner custom_ops_lib)
 
 # XNNPACK pthreadpool cpuinfo
 if(TARGET xnnpack_backend)
@@ -82,7 +91,8 @@ if(TARGET vulkan_backend)
   target_link_options_shared_lib(vulkan_backend)
 endif()
 
-target_compile_options(llama_main PUBLIC ${_common_compile_options})
+target_compile_options(llama_main PUBLIC ${_common_compile_options}
+  -DET_USE_THREADPOOL)
 target_link_libraries(llama_main PUBLIC ${link_libraries})
 
 if(APPLE)
diff --git a/examples/models/llama2/custom_ops/CMakeLists.txt b/examples/models/llama2/custom_ops/CMakeLists.txt
@@ -0,0 +1,79 @@
+# 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.
+
+
+cmake_minimum_required(VERSION 3.19)
+
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+if(NOT CMAKE_CXX_STANDARD)
+  set(CMAKE_CXX_STANDARD 17)
+endif()
+
+if(NOT PYTHON_EXECUTABLE)
+  set(PYTHON_EXECUTABLE python3)
+endif()
+
+# Source root directory for executorch.
+if(NOT EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..)
+endif()
+
+# Source root directory for pytorch.
+if(NOT TORCH_ROOT)
+  set(TORCH_ROOT ${EXECUTORCH_ROOT}/third-party/pytorch)
+endif()
+
+set(_common_compile_options -Wno-deprecated-declarations)
+
+include(${EXECUTORCH_ROOT}/build/Utils.cmake)
+include(${EXECUTORCH_ROOT}/build/Codegen.cmake)
+
+#
+# The `_<target>_srcs` lists are defined by including ${EXECUTORCH_SRCS_FILE}.
+#
+set(EXECUTORCH_SRCS_FILE
+    "${CMAKE_CURRENT_BINARY_DIR}/../../../../executorch_srcs.cmake")
+
+extract_sources(${EXECUTORCH_SRCS_FILE})
+
+include(${EXECUTORCH_SRCS_FILE})
+
+# Let files say "include <executorch/path/to/header.h>".
+set(_common_include_directories ${EXECUTORCH_ROOT}/..)
+
+
+# Custom op libraries
+set(custom_ops_libs extension_module)
+list(APPEND custom_ops_libs pthreadpool)
+list(APPEND custom_ops_libs cpuinfo)
+
+# Generate C++ bindings to register kernels into both PyTorch (for AOT) and
+# Executorch (for runtime). Here select all ops in optimized.yaml
+set(_yaml "${CMAKE_CURRENT_LIST_DIR}/custom_ops.yaml")
+gen_selected_ops("${_yaml}" "" "")
+
+generate_bindings_for_kernels(
+  FUNCTIONS_YAML ${CMAKE_CURRENT_SOURCE_DIR}/custom_ops.yaml)
+message("Generated files ${gen_command_sources}")
+
+list(TRANSFORM _custom_ops__srcs PREPEND "${EXECUTORCH_ROOT}/")
+add_library(custom_ops ${_custom_ops__srcs})
+
+target_include_directories(custom_ops PUBLIC "${_common_include_directories}")
+target_include_directories(custom_ops PRIVATE
+  "${CMAKE_CURRENT_BINARY_DIR}/../../../../include")
+target_link_libraries(custom_ops PRIVATE ${custom_ops_libs})
+
+target_compile_options(custom_ops PUBLIC ${_common_compile_options}
+  -DET_USE_THREADPOOL)
+
+# Build a library for _custom_ops_srcs
+#
+# custom_ops_lib: Register optimized ops kernels into Executorch runtime
+gen_operators_lib(
+  "custom_ops_lib"
+  KERNEL_LIBS custom_ops
+  DEPS executorch)
diff --git a/examples/models/llama2/runner/CMakeLists.txt b/examples/models/llama2/runner/CMakeLists.txt
@@ -47,10 +47,12 @@ else()
   add_library(llama_runner SHARED ${_llama_runner__srcs})
 endif()
 
-set(llama_runner_deps executorch extension_module extension_data_loader)
+set(llama_runner_deps executorch extension_module extension_data_loader
+  custom_ops)
 
 target_link_libraries(
   llama_runner PUBLIC ${llama_runner_deps})
 
 target_include_directories(llama_runner
-                           INTERFACE ${_common_include_directories} ${EXECUTORCH_ROOT})
+                           INTERFACE ${_common_include_directories}
+                           ${EXECUTORCH_ROOT})
