determine fht_avx.c/fht_neon.c based on CMake target processor

build/Utils.cmake

@@ -329,29 +329,25 @@ function(resolve_python_executable)
 endfunction()
 
 # find_package(Torch CONFIG REQUIRED) replacement for targets that have a
-# header-only Torch dependency. Because find_package sets variables in the
-# parent scope, we use a macro to preserve this rather than maintaining our own
-# list of those variables.
-macro(find_package_torch_headers)
-  # We cannot simply use CMAKE_FIND_ROOT_PATH_BOTH, because that does not
-  # propagate into TorchConfig.cmake.
-  foreach(mode_kind IN ITEMS PACKAGE LIBRARY INCLUDE)
-    set(OLD_CMAKE_FIND_ROOT_PATH_MODE_${mode_kind}
-        ${CMAKE_FIND_ROOT_PATH_MODE_${mode_kind}}
-    )
-    set(CMAKE_FIND_ROOT_PATH_MODE_${mode_kind} BOTH)
-  endforeach()
-  find_package_torch()
-  foreach(mode_kind IN ITEMS PACKAGE LIBRARY INCLUDE)
-    set(CMAKE_FIND_ROOT_PATH_MODE_${mode_kind}
-        ${OLD_CMAKE_FIND_ROOT_PATH_MODE_${mode_kind}}
-    )
-  endforeach()
-endmacro()
+# header-only Torch dependency.
+#
+# Unlike find_package(Torch ...), this will only set
+# TORCH_INCLUDE_DIRS in the parent scope. In particular, it will NOT
+# set any of the following:
+# - TORCH_FOUND
+# - TORCH_LIBRARY
+# - TORCH_CXX_FLAGS
+function(find_package_torch_headers)
+  # We implement this way rather than using find_package so that
+  # cross-compilation can still use the host's installed copy of
+  # torch, since the headers should be fine.
+  get_torch_base_path(TORCH_BASE_PATH)
+  set(TORCH_INCLUDE_DIRS "${TORCH_BASE_PATH}/include;${TORCH_BASE_PATH}/include/torch/csrc/api/include" PARENT_SCOPE)
+endfunction()
 
-# Add the Torch CMake configuration to CMAKE_PREFIX_PATH so that find_package
-# can find Torch.
-function(add_torch_to_cmake_prefix_path)
+# Return the base path to the installed Torch Python library in
+# outVar.
+function(get_torch_base_path outVar)
   if(NOT PYTHON_EXECUTABLE)
     resolve_python_executable()
   endif()
@@ -370,6 +366,13 @@ function(add_torch_to_cmake_prefix_path)
     message("Output:\n${_tmp_torch_path}")
     message(FATAL_ERROR "Error:\n${_tmp_torch_path_error}")
   endif()
+  set(${outVar} ${_tmp_torch_path} PARENT_SCOPE)
+endfunction()
+
+# Add the Torch CMake configuration to CMAKE_PREFIX_PATH so that find_package
+# can find Torch.
+function(add_torch_to_cmake_prefix_path)
+  get_torch_base_path(_tmp_torch_path)
   list(APPEND CMAKE_PREFIX_PATH "${_tmp_torch_path}")
   set(CMAKE_PREFIX_PATH
       "${CMAKE_PREFIX_PATH}"

build/build_android_llm_demo.sh

@@ -69,7 +69,7 @@ build_android_native_library() {
   fi
   cmake --build "${CMAKE_OUT}" -j "${CMAKE_JOBS}" --target install --config "${EXECUTORCH_CMAKE_BUILD_TYPE}"
 
-  cmake extension/android \
+  cmake --trace extension/android \
     -DCMAKE_TOOLCHAIN_FILE=${ANDROID_NDK}/build/cmake/android.toolchain.cmake \
     -DANDROID_ABI="${ANDROID_ABI}" \
     -DANDROID_PLATFORM=android-26 \

build/cmake_deps.toml

@@ -386,10 +386,7 @@ buck_targets = [
   "//extension/llm/custom_ops:custom_ops",
 ]
 filters = [
-  # Second clause is to pick up fht_neon.c/fht_avx.c from FFHT. TODO:
-  # remove filters and patch extract_sources.py's Buck query to fetch
-  # srcs; presumably filters is here to remove .h files.
-  "(.cpp$)|(fht.*\\.c$)",
+  ".cpp$",
 ]
 excludes = [
   "^codegen",

extension/llm/custom_ops/CMakeLists.txt

@@ -45,6 +45,22 @@ list(APPEND custom_ops_libs cpuinfo)
 list(APPEND custom_ops_libs cpublas)
 list(APPEND custom_ops_libs eigen_blas)
 
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64|armv7)$")
+  list(APPEND _custom_ops__srcs
+       "extension/llm/custom_ops/spinquant/third-party/FFHT/fht_neon.c"
+  )
+elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "x86_64")
+  list(APPEND _custom_ops__srcs
+       "extension/llm/custom_ops/spinquant/third-party/FFHT/fht_avx.c"
+  )
+else()
+  message(
+    FATAL_ERROR
+      "Unsupported CMAKE_SYSTEM_PROCESSOR ${CMAKE_SYSTEM_PROCESSOR}. (If \
+32-bit x86, try using fht_avx.c and send a PR if it works!)"
+  )
+endif()
+
 list(TRANSFORM _custom_ops__srcs PREPEND "${EXECUTORCH_ROOT}/")
 
 if(NOT EXECUTORCH_BUILD_XNNPACK)
@@ -121,3 +137,8 @@ if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
           LIBRARY DESTINATION executorch/extension/llm/custom_ops
   )
 endif()
+
+add_subdirectory(spinquant/third-party/FFHT)
+if(BUILD_TESTING)
+  add_subdirectory(spinquant/test)
+endif()

extension/llm/custom_ops/spinquant/test/CMakeLists.txt

@@ -0,0 +1,30 @@
+# 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.
+
+# @generated by test/utils/generate_gtest_cmakelists.py
+#
+# This file should be formatted with
+# ~~~
+# cmake-format -i CMakeLists.txt
+# ~~~
+# It should also be cmake-lint clean.
+#
+
+cmake_minimum_required(VERSION 3.19)
+
+set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..)
+
+include(${EXECUTORCH_ROOT}/build/Test.cmake)
+
+set(_test_srcs
+    fast_hadamard_transform_test.cpp fast_hadamard_transform_test_impl.cpp
+    op_fast_hadamard_transform_test.cpp
+)
+
+et_cxx_test(
+  extension_llm_custom_ops_spinquant_test SOURCES ${_test_srcs} EXTRA_LIBS
+  custom_ops dumb_fht
+)

extension/llm/custom_ops/spinquant/third-party/FFHT/CMakeLists.txt

@@ -0,0 +1,12 @@
+# 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.
+
+# Please this file formatted by running:
+# ~~~
+# cmake-format -i CMakeLists.txt
+# ~~~
+
+add_library(dumb_fht dumb_fht.c)

test/run_oss_cpp_tests.sh

@@ -35,6 +35,7 @@ build_executorch() {
   cmake . \
     -DCMAKE_INSTALL_PREFIX=cmake-out \
     -DEXECUTORCH_USE_CPP_CODE_COVERAGE=ON \
+    -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \
     -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \
     -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
     -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \

test/utils/OSSTestConfig.json

@@ -1,4 +1,15 @@
 { "tests": [
+    {
+        "directory": "extension/llm/custom_ops/spinquant/test",
+        "sources": [
+            "fast_hadamard_transform_test.cpp",
+            "fast_hadamard_transform_test_impl.cpp",
+            "op_fast_hadamard_transform_test.cpp"
+        ],
+        "additional_libs": [
+            "custom_ops"
+        ]
+    },
     {
         "directory": "extension/data_loader/test",
         "sources": [