Fix up CMakeLists and reorganize some code locations

torchao/experimental/CMakeLists.txt

@@ -0,0 +1,50 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+project(torchao)
+
+cmake_minimum_required(VERSION 3.19)
+
+set(CMAKE_CXX_STANDARD 17)
+
+if (NOT CMAKE_BUILD_TYPE)
+    set(CMAKE_BUILD_TYPE Release)
+endif()
+
+
+# Source root directory for torchao/experimental
+if(NOT TORCHAO_ROOT)
+  set(TORCHAO_ROOT ${CMAKE_CURRENT_SOURCE_DIR})
+endif()
+
+if(NOT TORCHAO_INCLUDE_DIRS)
+  set(TORCHAO_INCLUDE_DIRS ${TORCHAO_ROOT}/../..)
+endif()
+
+if (NOT TORCHAO_PARALLEL_BACKEND)
+    if (TORCHAO_OP_TARGET STREQUAL "ATEN")
+        set(TORCHAO_PARALLEL_BACKEND "ATEN_OPENMP")
+    elseif(TORCHAO_OP_TARGET STREQUAL "EXECUTORCH")
+        set(TORCHAO_PARALLEL_BACKEND "PTHREADPOOL")
+    else()
+        message(TORCHAO_PARALLEL_BACKEND "TORCHAO_PARALLEL_BACKEND is not set.  Please set it directly or set TORCHAO_OP_TARGET to get a default.")
+    endif()
+endif()
+
+include(CMakePrintHelpers)
+
+add_compile_options("-Wall" "-Werror")
+
+include(CMakePrintHelpers)
+message("TORCHAO_INCLUDE_DIRS: ${TORCHAO_INCLUDE_DIRS}")
+include_directories(${TORCHAO_INCLUDE_DIRS})
+
+if (CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
+    # Defines target torchao_kernels_aarch64
+    add_subdirectory(${TORCHAO_ROOT}/kernels/cpu/aarch64)
+    add_subdirectory(${TORCHAO_ROOT}/ops/linear)
+    add_subdirectory(${TORCHAO_ROOT}/ops/linear/linear_a8wxdq_op)
+endif()

torchao/experimental/kernels/cpu/linear/examples/torch_custom_op/build_custom_op.sh → torchao/experimental/build_torchao_ops.sh

@@ -5,15 +5,14 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)
-export TORCHAO_INCLUDE_DIRS=${SCRIPT_DIR}/../../../../../../..
-
 export CMAKE_PREFIX_PATH="$(python -c 'import torch.utils; print(torch.utils.cmake_prefix_path)')"
 echo "CMAKE_PREFIX_PATH: ${CMAKE_PREFIX_PATH}"
 export CMAKE_OUT=/tmp/cmake-out/torchao
-cmake -DTORCHAO_INCLUDE_DIRS=${TORCHAO_INCLUDE_DIRS} \
-    -DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH} \
-    -DPLATFORM="ATEN" \
-    -S ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/linear/examples/torch_custom_op \
+cmake -DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH} \
+    -DCMAKE_INSTALL_PREFIX=${CMAKE_OUT} \
+    -DTORCHAO_OP_TARGET="$1" \
+    -DEXECUTORCH_LIBRARIES=${EXECUTORCH_LIBRARIES} \
+    -DEXECUTORCH_INCLUDE_DIRS=${EXECUTORCH_INCLUDE_DIRS} \
+    -S . \
     -B ${CMAKE_OUT}
-cmake --build  ${CMAKE_OUT}
+cmake --build  ${CMAKE_OUT} --target install --config Release

torchao/experimental/kernels/cpu/aarch64/CMakeLists.txt

@@ -4,10 +4,12 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-add_library(
-  kernel_aarch64
-  ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/reduction/find_min_and_max.cpp
-  ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/reduction/compute_sum.cpp
-  ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/quantization/quantize.cpp
-  ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/valpacking/interleave.cpp
-)
+if (CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
+  add_library(
+    torchao_kernels_aarch64
+    ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/reduction/find_min_and_max.cpp
+    ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/reduction/compute_sum.cpp
+    ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/quantization/quantize.cpp
+    ${TORCHAO_INCLUDE_DIRS}/torchao/experimental/kernels/cpu/aarch64/valpacking/interleave.cpp
+  )
+endif()

torchao/experimental/ops/linear/CMakeLists.txt

@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+cmake_minimum_required(VERSION 3.19)
+
+include(${TORCHAO_ROOT}/Utils.cmake)
+
+add_library(torchao_ops_linear_${TORCHAO_PARALLEL_BACKEND} STATIC channelwise_8bit_activation_groupwise_lowbit_weight.cpp)
+target_link_torchao_parallel_backend(torchao_ops_linear_${TORCHAO_PARALLEL_BACKEND} "${TORCHAO_PARALLEL_BACKEND}")