Fix up CMakeLists and reorganize some code locations

Summary:
This diff cleans up CMakeFiles and reorganizes some of the code directories.

Archetecture-specific kernel code remains in experimental/kernels/cpu/aarch64.  This is not changed.

Code related to ops (parallel.h, memory.h, "high-level" op interface, examples/torch_custom_op, etc) is moved to experimental/ops.

The example code that builds the custom ops for ATen/ExecuTorch (experimental/kernels/cpu/linear/examples/torch_custom_op) is moved out of examples and into a more descriptive directory: experimental/ops/linear/linear_a8wxdq_op.  These are the kernels that will be used in torchchat initially.

The "high-level" op interface is in ops/linear/channelwise_8bit_activation_groupwise_lowbit_weight.h.  A library for it is defined in the CMakeFiles.txt.  It uses no templates and is not specific to the universal kernels and can be used with KleidiAI kernels.  The ops in experimental/ops/linear/linear_a8wxdq_op use the high-level library and the universal kernels in  experimental/kernels/cpu/aarch64 to define custom ops.

Finally, there is a CMakeLists at the top-level directory experimental.  The script build_torchao_ops.sh builds the torchao_ops for different platforms:

```
sh build_torchao_ops.sh ATEN
```

For ExecuTorch, you must define the environment variables EXECUTORCH_INCLUDE_DIRS and EXECUTORCH_LIBRARIES manually before calling:

```
EXECUTORCH_INCLUDE_DIRS=$HOME EXECUTORCH_LIBRARIES=$HOME/executorch/cmake-out/lib/libexecutorch_no_prim_ops.a sh build_torchao_ops.sh EXECUTORCH
```

Reviewed By: digantdesai

Differential Revision: D62711903

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}")