AQLM custom kernels for Android

CMakeLists.txt

@@ -467,6 +467,9 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM)
   add_subdirectory(
     ${CMAKE_CURRENT_SOURCE_DIR}/examples/models/llama2/custom_ops
   )
+  add_subdirectory(
+    ${CMAKE_CURRENT_SOURCE_DIR}/examples/models/llama2/aqlm
+  )
 endif()
 
 if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
@@ -633,13 +636,16 @@ if(EXECUTORCH_BUILD_PYBIND)
   # TODO(larryliu): Fix macOS 2 dylibs having 2 sets of static variables issue
   if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT AND NOT APPLE)
     list(APPEND _dep_libs custom_ops_aot_lib)
+    list(APPEND _dep_libs aqlm_aot_lib)
   endif()
   # TODO(laryliu): Fix linux duplicate registation problem. In GH CI worker
   # libcustom_ops.a doesn't dedup with the one indirectly linked from
   # libcustom_ops_aot_lib.a
   if(EXECUTORCH_BUILD_KERNELS_CUSTOM AND APPLE)
     target_link_options_shared_lib(custom_ops)
     list(APPEND _dep_libs custom_ops)
+    target_link_options_shared_lib(aqlm)
+    list(APPEND _dep_libs aqlm)
   endif()
   # compile options for pybind
   set(_pybind_compile_options
@@ -699,7 +705,7 @@ if(EXECUTORCH_BUILD_PYBIND)
       PROPERTIES # Assume that this library will be installed in
                  # `site-packages/executorch/extension/pybindings`, and that
                  # the custom_ops_aot_lib should be found with relative path.
-                 BUILD_RPATH "$ORIGIN:$ORIGIN/../../examples/models/llama2/custom_ops"
+                 BUILD_RPATH "$ORIGIN:$ORIGIN/../../examples/models/llama2/custom_ops:$ORIGIN/../../examples/models/llama2/aqlm"
     )
   endif()
 

examples/demo-apps/android/LlamaDemo/app/src/main/java/com/example/executorchllamademo/MainActivity.java

@@ -202,6 +202,9 @@ private void onModelRunStopped() {
     mSendButton.setOnClickListener(
         view -> {
           String prompt = mEditTextMessage.getText().toString();
+
+          String chat_prompt = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n" + prompt + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n";
+
           mMessageAdapter.add(new Message(prompt, true));
           mMessageAdapter.notifyDataSetChanged();
           mEditTextMessage.setText("");
@@ -219,7 +222,7 @@ public void run() {
                         }
                       });
 
-                  mModule.generate(prompt, MainActivity.this);
+                  mModule.generate(chat_prompt, MainActivity.this);
 
                   runOnUiThread(
                       new Runnable() {

examples/models/llama2/CMakeLists.txt

@@ -87,6 +87,7 @@ endif()
 # custom ops library
 if(EXECUTORCH_BUILD_KERNELS_CUSTOM)
   add_subdirectory(custom_ops)
+  add_subdirectory(aqlm)
 endif()
 
 # llama_runner library
@@ -129,6 +130,9 @@ list(APPEND link_libraries quantized_kernels quantized_ops_lib)
 if(EXECUTORCH_BUILD_KERNELS_CUSTOM)
   target_link_options_shared_lib(custom_ops)
   list(APPEND link_libraries custom_ops)
+
+  target_link_options_shared_lib(aqlm)
+  list(APPEND link_libraries aqlm)
 endif()
 
 set(XNNPACK_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../backends/xnnpack)

examples/models/llama2/aqlm/CmakeLists.txt

@@ -0,0 +1,111 @@
+# 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()
+
+set(_common_compile_options -Wno-deprecated-declarations -fPIC)
+
+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(aqlm_libs executorch_no_prim_ops)
+list(APPEND aqlm_libs pthreadpool)
+list(APPEND aqlm_libs cpuinfo)
+list(APPEND aqlm_libs cpublas)
+list(APPEND aqlm_libs eigen_blas)
+
+set(_aqlm__srcs examples/models/llama2/aqlm/lut_kernel.h examples/models/llama2/aqlm/lut_kernel.cpp)
+list(TRANSFORM _aqlm__srcs PREPEND "${EXECUTORCH_ROOT}/")
+
+
+
+message("HERE: AQLM SOURCES: ${_aqlm__srcs}")
+
+# TODO: Consider moving xnnpack/threadpool in a separate lib since it's now used
+# by custom ops too.
+if(NOT EXECUTORCH_BUILD_XNNPACK)
+  list(
+    APPEND
+    _aqlm__srcs
+    "${CMAKE_CURRENT_SOURCE_DIR}/../../../../backends/xnnpack/threadpool/threadpool.cpp"
+    "${CMAKE_CURRENT_SOURCE_DIR}/../../../../backends/xnnpack/threadpool/threadpool_guard.cpp"
+  )
+else()
+  list(APPEND aqlm_libs xnnpack_backend)
+endif()
+
+find_package(OpenMP REQUIRED)
+# list(APPEND aqlm_libs OpenMP::OpenMP_CXX)
+list(APPEND aqlm_libs omp)
+
+add_library(aqlm ${_aqlm__srcs})
+
+# Enable optimization
+set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
+
+target_include_directories(aqlm PUBLIC "${_common_include_directories}")
+target_include_directories(
+  aqlm PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/../../../../include"
+)
+target_link_libraries(aqlm PUBLIC ${aqlm_libs} -fopenmp -static-openmp)
+
+target_compile_options(
+  aqlm PUBLIC ${_common_compile_options} -DET_USE_THREADPOOL
+)
+
+install(TARGETS aqlm DESTINATION lib)
+
+if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
+  # Add a AOT library
+  find_package(Torch CONFIG REQUIRED)
+  add_library(
+    aqlm_aot_lib SHARED ${CMAKE_CURRENT_SOURCE_DIR}/lut_kernel_pytorch.cpp
+  )
+  target_include_directories(
+    aqlm_aot_lib PUBLIC "${_common_include_directories}"
+  )
+  target_include_directories(
+    aqlm_aot_lib
+    PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/../../../../include"
+  )
+  target_link_libraries(aqlm_aot_lib PUBLIC aqlm torch)
+  target_compile_options(
+    aqlm_aot_lib PUBLIC -Wno-deprecated-declarations -fPIC -frtti
+                              -fexceptions
+  )
+
+  install(TARGETS aqlm_aot_lib DESTINATION lib)
+endif()