cmake torchao_ops_mps_linear_fp_act_xbit_weight

Summary: Move from setup.py to cmake for building custom torchao mps ops

Differential Revision: D66120124

Author: manuelcandales

torchao/experimental/ops/mps/CMakeLists.txt

@@ -0,0 +1,45 @@
+# 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)
+
+project(torchao_ops_mps_linear_fp_act_xbit_weight)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED YES)
+
+if (NOT CMAKE_BUILD_TYPE)
+    set(CMAKE_BUILD_TYPE Release)
+endif()
+
+find_package(Torch REQUIRED)
+
+if(NOT TORCHAO_INCLUDE_DIRS)
+  set(TORCHAO_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/../../../..)
+endif()
+message(STATUS "TORCHAO_INCLUDE_DIRS: ${TORCHAO_INCLUDE_DIRS}")
+
+include_directories(${TORCHAO_INCLUDE_DIRS})
+add_library(torchao_ops_mps_linear_fp_act_xbit_weight_aten SHARED aten/register.mm)
+
+target_include_directories(torchao_ops_mps_linear_fp_act_xbit_weight_aten PRIVATE "${TORCH_INCLUDE_DIRS}")
+target_link_libraries(torchao_ops_mps_linear_fp_act_xbit_weight_aten PRIVATE "${TORCH_LIBRARIES}")
+target_compile_definitions(torchao_ops_mps_linear_fp_act_xbit_weight_aten PRIVATE USE_ATEN=1)
+
+# Enable Metal support
+if (APPLE)
+    find_library(METAL_LIB Metal)
+    find_library(FOUNDATION_LIB Foundation)
+    target_link_libraries(torchao_ops_mps_linear_fp_act_xbit_weight_aten PRIVATE ${METAL_LIB} ${FOUNDATION_LIB})
+else()
+    message(FATAL_ERROR "Torchao experimental mps ops can only be built on macOS/iOS")
+endif()
+
+install(
+    TARGETS torchao_ops_mps_linear_fp_act_xbit_weight_aten
+    EXPORT _targets
+    DESTINATION lib
+)

torchao/experimental/ops/mps/register.mm renamed to torchao/experimental/ops/mps/aten/register.mm

@@ -5,7 +5,7 @@
 // LICENSE file in the root directory of this source tree.
 
 // clang-format off
-#include <torch/extension.h>
+#include <torch/library.h>
 #include <ATen/native/mps/OperationUtils.h>
 #include <torchao/experimental/kernels/mps/src/lowbit.h>
 // clang-format on
@@ -147,9 +147,6 @@ Tensor pack_weights_cpu_kernel(const Tensor& W) {
   return B;
 }
 
-// Registers _C as a Python extension module.
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {}
-
 TORCH_LIBRARY(torchao, m) {
   m.def("_pack_weight_1bit(Tensor W) -> Tensor");
   m.def("_pack_weight_2bit(Tensor W) -> Tensor");

torchao/experimental/ops/mps/build.sh

@@ -0,0 +1,20 @@
+#!/bin/bash -eu
+# 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.
+
+python ../../kernels/mps/codegen/gen_metal_shader_lib.py
+
+export CMAKE_PREFIX_PATH=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')
+echo "CMAKE_PREFIX_PATH: ${CMAKE_PREFIX_PATH}"
+export CMAKE_OUT=cmake-out
+
+cmake -DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH} \
+    -DCMAKE_INSTALL_PREFIX=${CMAKE_OUT} \
+    -S . \
+    -B ${CMAKE_OUT}
+cmake --build  ${CMAKE_OUT} -j 16 --target install --config Release
+
+rm ../../kernels/mps/src/metal_shader_lib.h

torchao/experimental/ops/mps/setup.py

@@ -4,10 +4,17 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
+import os
 import torch
-import torchao_mps_ops
 import unittest
 
+path_libtorchao_ops_mps_aten = os.path.abspath(
+    os.path.join(
+        os.path.dirname(__file__), "../cmake-out/lib/libtorchao_ops_mps_linear_fp_act_xbit_weight_aten.dylib"
+    )
+)
+torch.ops.load_library(path_libtorchao_ops_mps_aten)
+
 
 def parameterized(test_cases):
     def decorator(func):