Update on "Reuse GELU implementation from PyTorch core"

kernels/optimized doesn't need to support embedded systems, so it can just take a header-only dep on PyTorch.

Note that, because we will pick up Sleef internally and ignore it
externally thanks to ATen vec, this PR gets to enable optimized GELU in OSS.

Testing: CI to make sure this doesn't break mobile build modes; happy to take advice on anything not currently covered that might break.

Differential Revision: [D66335522](https://our.internmc.facebook.com/intern/diff/D66335522/)

[ghstack-poisoned]

Author: swolchok

.ci/scripts/build_llama_android.sh

@@ -15,7 +15,6 @@ if [[ -z "${PYTHON_EXECUTABLE:-}" ]]; then
 fi
 which "${PYTHON_EXECUTABLE}"
 CMAKE_PREFIX_PATH="$(python3 -c 'import torch as _; print(_.__path__[0])')"
-SITE_PACKAGES="$(${PYTHON_EXECUTABLE} -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')"
 
 install_executorch_and_backend_lib() {
   echo "Installing executorch and xnnpack backend"
@@ -29,7 +28,6 @@ install_executorch_and_backend_lib() {
     -DANDROID_ABI="${ANDROID_ABI}" \
     -DCMAKE_INSTALL_PREFIX=cmake-android-out \
     -DCMAKE_BUILD_TYPE=Release \
-    -DCMAKE_FIND_ROOT_PATH="${SITE_PACKAGES}" \
     -DCMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH}" \
     -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
     -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \

build/build_android_llm_demo.sh

@@ -40,10 +40,8 @@ build_android_native_library() {
   else
     EXECUTORCH_BUILD_NEURON=OFF
   fi
-  SITE_PACKAGES="$(${PYTHON_EXECUTABLE} -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')"
 
   cmake . -DCMAKE_INSTALL_PREFIX="${CMAKE_OUT}" \
-    -DCMAKE_FIND_ROOT_PATH="${SITE_PACKAGES}" \
     -DCMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH}" \
     -DCMAKE_TOOLCHAIN_FILE="${ANDROID_NDK}/build/cmake/android.toolchain.cmake" \
     -DANDROID_ABI="${ANDROID_ABI}" \
@@ -78,7 +76,6 @@ build_android_native_library() {
     -DANDROID_ABI="${ANDROID_ABI}" \
     -DANDROID_PLATFORM=android-26 \
     -DCMAKE_INSTALL_PREFIX="${CMAKE_OUT}" \
-    -DCMAKE_FIND_ROOT_PATH="${SITE_PACKAGES}" \
     -DCMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH}" \
     -DEXECUTORCH_ENABLE_LOGGING=ON \
     -DEXECUTORCH_LOG_LEVEL=Info \

kernels/optimized/CMakeLists.txt

@@ -63,7 +63,14 @@ message("Generated files ${gen_command_sources}")
 
 list(TRANSFORM _optimized_kernels__srcs PREPEND "${EXECUTORCH_ROOT}/")
 add_library(optimized_kernels ${_optimized_kernels__srcs})
-find_package(Torch CONFIG REQUIRED)
+# We require Torch headers, which setup.py puts in CMAKE_PREFIX_PATH
+# for us. Toolchains that we might be using for cross-compiling could
+# set CMAKE_FIND_ROOT_PATH, which prevents find_package from finding
+# headers not rooted under CMAKE_FIND_ROOT_PATH. This is reasonable
+# for binary dependencies because they probably aren't built for the
+# target platform, but for our header-only use case, we should just
+# ignore CMAKE_FIND_ROOT_PATH.
+find_package(Torch CONFIG REQUIRED NO_CMAKE_FIND_ROOT_PATH)
 target_include_directories(optimized_kernels PRIVATE ${TORCH_INCLUDE_DIRS})
 target_link_libraries(
   optimized_kernels PRIVATE executorch_core cpublas extension_threadpool