Add ENABLE_KLEIDI_FLAG to use new QP8 Kernels

pytorch · Sep 14, 2024 · 4d96001 · 4d96001
1 parent 61a8270
commit 4d96001
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 7 deletions.
diff --git a/backends/xnnpack/CMakeLists.txt b/backends/xnnpack/CMakeLists.txt
@@ -32,14 +32,17 @@ if(NOT PYTHON_EXECUTABLE)
   resolve_python_executable()
 endif()
 
-# NB: Enabling this will serialize execution of delegate instances Keeping this
-# OFF by default to maintain existing behavior, to be revisited.
-option(EXECUTORCH_XNNPACK_SHARED_WORKSPACE
-       "Enable workspace sharing across different delegate instances" OFF
-)
+# NB: Enabling this will serialize execution of delegate instances
+# Keeping this OFF by default to maintain existing behavior, to be revisited.
+option(EXECUTORCH_XNNPACK_SHARED_WORKSPACE "Enable workspace sharing across different delegate instances" OFF)
+# Keeping this OFF by default do to regressions in decode and model load with kleidi kernels
+option(EXECUTORCH_XNNPACK_ENABLE_KLEIDI "Enable workspace sharing across different delegate instances" OFF)
 if(EXECUTORCH_XNNPACK_SHARED_WORKSPACE)
   add_definitions(-DENABLE_XNNPACK_SHARED_WORKSPACE)
 endif()
+if(EXECUTORCH_XNNPACK_ENABLE_KLEIDI)
+  add_definitions(-DENABLE_XNNPACK_KLEIDI)
+endif()
 
 set(_common_include_directories ${EXECUTORCH_ROOT}/..)
 set(_common_compile_options -Wno-deprecated-declarations -fPIC)

diff --git a/backends/xnnpack/cmake/Dependencies.cmake b/backends/xnnpack/cmake/Dependencies.cmake
@@ -36,7 +36,21 @@ set(XNNPACK_ENABLE_AVXVNNI
     OFF
     CACHE BOOL ""
 )
-set(XNNPACK_ENABLE_KLEIDIAI
+
+if(EXECUTORCH_XNNPACK_ENABLE_KLEIDI)
+    set(XNNPACK_ENABLE_KLEIDIAI
+        ON
+        CACHE BOOL ""
+    )
+else()
+    set(XNNPACK_ENABLE_KLEIDIAI
+        OFF
+        CACHE BOOL ""
+    )
+endif()
+
+
+set(XNNPACK_BUILD_ALL_MICROKERNELS
     OFF
     CACHE BOOL ""
 )
@@ -48,6 +62,14 @@ install(TARGETS microkernels-prod
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     PUBLIC_HEADER DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
 
+
+if(EXECUTORCH_XNNPACK_ENABLE_KLEIDI)
+    install(TARGETS kleidiai
+        LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+        ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+        PUBLIC_HEADER DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
+endif()
+
 # Revert PIC Flag to what it originally was
 set(CMAKE_POSITION_INDEPENDENT_CODE
     ${ORIGINAL_CMAKE_POSITION_INDEPENDENT_CODE_FLAG}

diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp
@@ -630,7 +630,14 @@ Error defineConvertNode(
       subgraph_ptr,
       remapped_ids.at(graph_node->input_id()),
       remapped_ids.at(graph_node->output_id()),
+#ifdef ENABLE_XNNPACK_KLEIDI
+      // This maps to XNNPACK's XNN_FLAG_MAYBE_PACK_FOR_GEMM
+      // however this is not currently exposed at top level
+      // Header
+      0x00000080);
+#else
       graph_node->flags());
+#endif
 
   ET_CHECK_OR_RETURN_ERROR(
       status == xnn_status_success,

diff --git a/backends/xnnpack/targets.bzl b/backends/xnnpack/targets.bzl
@@ -40,6 +40,7 @@ def define_common_targets():
             # "-DENABLE_XNNPACK_PROFILING",
             # Uncomment to enable workspace sharing across delegates
             # "-DENABLE_XNNPACK_SHARED_WORKSPACE"
+            # "-DENABLE_XNNPACK_KLEIDI"
         ],
         exported_deps = [
             "//executorch/runtime/backend:interface",

diff --git a/build/executorch-config.cmake b/build/executorch-config.cmake
@@ -49,8 +49,11 @@ set(lib_list
     extension_tensor
     extension_threadpool
     xnnpack_backend
+    # Start XNNPACK Lib Deps
     XNNPACK
     microkernels-prod
+    kleidiai
+    # End XNNPACK Lib Deps
     cpuinfo
     pthreadpool
     vulkan_backend

diff --git a/examples/models/llama2/CMakeLists.txt b/examples/models/llama2/CMakeLists.txt
@@ -142,6 +142,9 @@ endif()
 # XNNPACK
 if(TARGET xnnpack_backend)
   set(xnnpack_backend_libs xnnpack_backend XNNPACK microkernels-prod)
+  if(TARGET kleidiai)
+    list(APPEND xnnpack_backend_libs kleidiai)
+  endif()
   list(APPEND link_libraries ${xnnpack_backend_libs})
   target_link_options_shared_lib(xnnpack_backend)
 endif()

diff --git a/examples/models/llava/CMakeLists.txt b/examples/models/llava/CMakeLists.txt
@@ -157,7 +157,10 @@ endif()
 
 # XNNPACK
 if(TARGET xnnpack_backend)
-  set(xnnpack_backend_libs xnnpack_backend XNNPACK)
+  set(xnnpack_backend_libs xnnpack_backend XNNPACK microkernels-prod)
+  if(TARGET kleidiai)
+    list(APPEND xnnpack_backend_libs kleidiai)
+  endif()
   list(APPEND link_libraries ${xnnpack_backend_libs})
   target_link_options_shared_lib(xnnpack_backend)
 endif()