update torchao library name

Author: metascroy

runner/aoti.cmake

@@ -30,5 +30,5 @@ if(Torch_FOUND)
 endif()
 
 if (LINK_TORCHAO_CUSTOM_OPS)
-    target_link_libraries(aoti_run "${TORCHCHAT_ROOT}/torchao-build/cmake-out/liblowbit_op_aten${CMAKE_SHARED_LIBRARY_SUFFIX}")
+    target_link_libraries(aoti_run "${TORCHCHAT_ROOT}/torchao-build/cmake-out/ops/linear/linear_a8wxdq_op/liblinear_a8wxdq_aten${CMAKE_SHARED_LIBRARY_SUFFIX}")
 endif()

runner/et.cmake

@@ -112,7 +112,7 @@ if(executorch_FOUND)
   endif()
 
   if(LINK_TORCHAO_CUSTOM_OPS)
-    target_link_libraries(et_run PRIVATE "${TORCHCHAT_ROOT}/torchao-build/cmake-out/liblowbit_op_executorch${CMAKE_SHARED_LIBRARY_SUFFIX}")
+    target_link_libraries(et_run PRIVATE "${TORCHCHAT_ROOT}/torchao-build/cmake-out/ops/linear/linear_a8wxdq_op/liblinear_a8wxdq_aten${CMAKE_SHARED_LIBRARY_SUFFIX}")
   endif()
 
   # Adding target_link_options_shared_lib as commented out below leads to this:

torchchat/utils/quantize.py

@@ -891,7 +891,7 @@ def quantized_model(self) -> nn.Module:
     # Try loading custom op
     try:
         import glob
-        libs = glob.glob(f"{torchao_build_path}/cmake-out/liblowbit_op_aten.*")
+        libs = glob.glob(f"{torchao_build_path}/cmake-out/ops/linear/linear_a8wxdq_op/liblinear_a8wxdq_aten.*")
         libs = list(filter(lambda l: (l.endswith("so") or l.endswith("dylib")), libs))
         torch.ops.load_library(libs[0])
     except Exception as e:

torchchat/utils/scripts/install_utils.sh

@@ -180,30 +180,31 @@ clone_torchao() {
 
 install_torchao_custom_aten_ops() {
   echo "Building torchao custom ops for ATen"
-  pushd ${TORCHCHAT_ROOT}/torchao-build/src/ao/torchao/experimental/kernels/cpu/linear/examples/torch_custom_op
-  export TORCHAO_INCLUDE_DIRS=${TORCHCHAT_ROOT}/torchao-build/src/ao
+  pushd ${TORCHCHAT_ROOT}/torchao-build/src/ao/torchao/experimental
 
   CMAKE_OUT_DIR=${TORCHCHAT_ROOT}/torchao-build/cmake-out
-  cmake -DTORCHAO_INCLUDE_DIRS=${TORCHAO_INCLUDE_DIRS} \
-    -DCMAKE_PREFIX_PATH=${MY_CMAKE_PREFIX_PATH} \
-    -DPLATFORM="ATEN" \
+  cmake -DCMAKE_PREFIX_PATH=${MY_CMAKE_PREFIX_PATH} \
+    -DTORCHAO_OP_TARGET="ATEN" \
     -S . \
     -B ${CMAKE_OUT_DIR} -G Ninja
   cmake --build  ${CMAKE_OUT_DIR}
+
+  popd
 }
 
 install_torchao_custom_executorch_ops() {
   echo "Building torchao custom ops for ExecuTorch"
-  pushd ${TORCHCHAT_ROOT}/torchao-build/src/ao/torchao/experimental/kernels/cpu/linear/examples/torch_custom_op
+  pushd ${TORCHCHAT_ROOT}/torchao-build/src/ao/torchao/experimental
   export TORCHAO_INCLUDE_DIRS=${TORCHCHAT_ROOT}/torchao-build/src/ao
 
   CMAKE_OUT_DIR="${TORCHCHAT_ROOT}/torchao-build/cmake-out"
-  cmake -DTORCHAO_INCLUDE_DIRS=${TORCHAO_INCLUDE_DIRS} \
-    -DCMAKE_PREFIX_PATH=${MY_CMAKE_PREFIX_PATH} \
+  cmake -DCMAKE_PREFIX_PATH=${MY_CMAKE_PREFIX_PATH} \
     -DEXECUTORCH_INCLUDE_DIRS=${EXECUTORCH_INCLUDE_DIRS} \
     -DEXECUTORCH_LIBRARIES=${EXECUTORCH_LIBRARIES} \
-    -DPLATFORM="EXECUTORCH" \
+    -DTORCHAO_OP_TARGET="EXECUTORCH" \
     -S . \
     -B ${CMAKE_OUT_DIR} -G Ninja
   cmake --build  ${CMAKE_OUT_DIR}
+
+  popd
 }