Provide CMake-driver for building libcustom_calls

CMakeLists.txt

@@ -0,0 +1,13 @@
+cmake_minimum_required(VERSION 3.20)
+project(catalyst)
+
+set(LOGO [=[
+░█▀▀░█▀█░▀█▀░█▀█░█░░░█░█░█▀▀░▀█▀
+░█░░░█▀█░░█░░█▀█░█░░░░█░░▀▀█░░█░
+░▀▀▀░▀░▀░░▀░░▀░▀░▀▀▀░░▀░░▀▀▀░░▀░
+]=])
+message(${LOGO})
+
+add_subdirectory(frontend)
+add_subdirectory(mlir)
+add_subdirectory(runtime)

frontend/catalyst/CMakeLists.txt

@@ -1 +1,3 @@
+project(catalyst_frontend)
+
 add_subdirectory(utils)

frontend/catalyst/utils/CMakeLists.txt

@@ -1,9 +1,27 @@
+project(catalyst_frontend_utils)
+
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
+# Find the Python interpreter, development headers and NumPy headers
 find_package(Python 3
-    REQUIRED COMPONENTS Interpreter Development.Module
+    REQUIRED COMPONENTS Interpreter Development.Module NumPy
     OPTIONAL_COMPONENTS Development.SABIModule)
+include_directories(${PYTHON_INCLUDE_DIRS})
+
+# Build the custom calls library and make `catalyst_jax_cpu_lapack_kernels` 
+# target visible in the current project directory.
+add_subdirectory(jax_cpu_lapack_kernels)
+
+add_library(custom_calls SHARED
+    ${PROJECT_SOURCE_DIR}/libcustom_calls.cpp
+)
+target_link_libraries(custom_calls PUBLIC "$<LINK_LIBRARY:WHOLE_ARCHIVE,catalyst_jax_cpu_lapack_kernels>")
+set_property(TARGET custom_calls PROPERTY ENABLE_EXPORTS OFF)
+
+# If aiming to build additional shared submodules, setting the RPATH $ORIGIN
+# value may be useful to avoid explicit library paths.
+set_target_properties(custom_calls PROPERTIES BUILD_RPATH "$ORIGIN/")
 
 # nanobind suggests including these lines to configure CMake to perform an optimized release build
 # by default unless another build type is specified. Without this addition, binding code may run
@@ -21,13 +39,6 @@ execute_process(
 
 find_package(nanobind CONFIG REQUIRED)
 
-# Get the NumPy include directory
-execute_process(
-    COMMAND "${Python_EXECUTABLE}" -c "import numpy; print(numpy.get_include())"
-    OUTPUT_VARIABLE NUMPY_INCLUDE_DIR
-    OUTPUT_STRIP_TRAILING_WHITESPACE
-)
-
 # Source file list for `wrapper` module
 set(WRAPPER_SRC_FILES
     wrapper.cpp
@@ -39,8 +50,12 @@ set(WRAPPER_SRC_FILES
 nanobind_add_module(wrapper STABLE_ABI ${WRAPPER_SRC_FILES})
 
 # Add the NumPy include directory to the library's include paths
-target_include_directories(wrapper PRIVATE ${NUMPY_INCLUDE_DIR})
+target_include_directories(wrapper PRIVATE ${Python_NumPy_INCLUDE_DIRS})
 
 # Use suffix ".so" rather than ".abi3.so" for library file using Stable ABI
 # This is necessary for compatibility with setuptools build extensions
 set_target_properties(wrapper PROPERTIES SUFFIX ".so")
+
+# Allow custom calls to be loaded alongside wrapper
+target_link_libraries(wrapper INTERFACE custom_calls)
+set_target_properties(wrapper PROPERTIES BUILD_RPATH "$ORIGIN/")

frontend/catalyst/utils/jax_cpu_lapack_kernels/CMakeLists.txt

@@ -0,0 +1,36 @@
+project(catalyst_jax_cpu_lapack_kernels)
+
+# Create a static library for LAPACK kernels to be included into the
+# custom_calls librar in the parent scope.
+
+add_library(catalyst_jax_cpu_lapack_kernels STATIC
+    ${PROJECT_SOURCE_DIR}/lapack_kernels.cpp
+    ${PROJECT_SOURCE_DIR}/lapack_kernels_using_lapack.cpp
+)
+
+# Ensure fPIC is set for static builds
+set_property(TARGET catalyst_jax_cpu_lapack_kernels PROPERTY POSITION_INDEPENDENT_CODE ON)
+
+# Headers are tracked into the target from the current project scope 
+target_include_directories(catalyst_jax_cpu_lapack_kernels PUBLIC
+  $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}>
+  $<INSTALL_INTERFACE:include> 
+)
+
+# If the option is set, allow explicit linking of the CMake-found CBLAS/LAPACKE
+# library. See https://cmake.org/cmake/help/latest/module/FindBLAS.html for info.
+option(CATALYST_LINK_BLAS "Link the CMake-defined BLAS and LAPACK libraries as part of the extension module" OFF)
+
+if(CATALYST_LINK_BLAS)
+    find_package( BLAS REQUIRED )
+    find_package( LAPACK REQUIRED)
+    target_link_libraries(catalyst_jax_cpu_lapack_kernels PUBLIC BLAS::BLAS)
+
+    # Ensure libraries are explicitly linked.
+    target_link_options(catalyst_jax_cpu_lapack_kernels PUBLIC
+        $<$<COMPILE_LANGUAGE:CXX>:-Wl,--no-as-needed>
+    )
+endif()
+
+# Ensure the availability of symbol exports. Only necessary with `SHARED`
+set_property(TARGET catalyst_jax_cpu_lapack_kernels PROPERTY ENABLE_EXPORTS ON)

setup.py

@@ -16,7 +16,6 @@
 
 # pylint: disable=wrong-import-order
 
-import glob
 import os
 import platform
 import subprocess
@@ -126,30 +125,29 @@
 
 
 class CMakeExtension(Extension):
-    """A setuptools Extension class for modules with a CMake configuration."""
+    """setuptools Extension wrapper for CMake.
+
+    Provides access to the source directories directly for modules to be compiled.
+
+    For example, to build the `libcustom_calls` library, the direct module path can be provided as
+    ```python
+    CMakeExtension("catalyst.utils.libcustom_calls", sourcedir=frontend_dir)
+    ```
+    """
 
     def __init__(self, name, sourcedir=""):
-        super().__init__(name, sources=[])
+        Extension.__init__(self, name, sources=[])
         self.sourcedir = os.path.abspath(sourcedir)
 
 
-class UnifiedBuildExt(build_ext):
+class CMakeBuild(build_ext):
     """Custom build extension class for the Catalyst Frontend.
 
     This class overrides a number of methods from its parent class
     setuptools.command.build_ext.build_ext, the most important of which are:
 
         1. `get_ext_filename`, in order to remove the architecture/python
            version suffix of the library name.
-        2. `build_extension`, in order to handle the compilation of extensions
-           with CMake configurations, namely the catalyst.utils.wrapper module,
-           and of generic C/C++ extensions without a CMake configuration, namely
-           the catalyst.utils.libcustom_calls module, which is currently built
-           as a plain setuptools Extension.
-
-    TODO: Eventually it would be better to build the utils.libcustom_calls
-    module using a CMake configuration as well, rather than as a setuptools
-    Extension.
     """
 
     def initialize_options(self):
@@ -175,12 +173,6 @@ def get_ext_filename(self, fullname):
         return filename.replace(suffix, "") + extension
 
     def build_extension(self, ext):
-        if isinstance(ext, CMakeExtension):
-            self.build_cmake_extension(ext)
-        else:
-            super().build_extension(ext)
-
-    def build_cmake_extension(self, ext: CMakeExtension):
         """Configure and build CMake extension."""
         cmake_path = "cmake"
         ninja_path = "ninja"
@@ -230,81 +222,25 @@ def build_cmake_extension(self, ext: CMakeExtension):
         subprocess.check_call([cmake_path, "--build", "."] + build_args, cwd=build_temp)
 
 
-class CustomBuildExtLinux(UnifiedBuildExt):
-    """Custom build extension class for Linux platforms
-
-    Currently no extra work needs to be performed with respect to the base class
-    UnifiedBuildExt.
-    """
-
-
-class CustomBuildExtMacos(UnifiedBuildExt):
-    """Custom build extension class for macOS platforms
-
-    In addition to the work performed by the base class UnifiedBuildExt, this
-    class also changes the LC_ID_DYLIB that is otherwise constant and equal to
-    where the shared library was created.
-    """
-
-    def run(self):
-        # Run the original build_ext command
-        super().run()
-
-        # Construct library name based on ext suffix (contains python version, architecture and .so)
-        library_name = "libcustom_calls.so"
-
-        package_root = os.path.dirname(__file__)
-        frontend_path = glob.glob(
-            os.path.join(package_root, "frontend", "**", library_name), recursive=True
-        )
-        build_path = glob.glob(os.path.join("build", "**", library_name), recursive=True)
-        lib_with_r_path = "@rpath/libcustom_calls.so"
-
-        original_path = frontend_path[0] if frontend_path else build_path[0]
-
-        # Run install_name_tool to modify LC_ID_DYLIB(other the rpath stays in vars/folder)
-        subprocess.run(
-            ["/usr/bin/install_name_tool", "-id", lib_with_r_path, original_path],
-            check=False,
-        )
-
-
 # Compile the library of custom calls in the frontend
 if system_platform == "Linux":
-    custom_calls_extension = Extension(
-        "catalyst.utils.libcustom_calls",
-        sources=[
-            "frontend/catalyst/utils/libcustom_calls.cpp",
-            "frontend/catalyst/utils/jax_cpu_lapack_kernels/lapack_kernels.cpp",
-            "frontend/catalyst/utils/jax_cpu_lapack_kernels/lapack_kernels_using_lapack.cpp",
-        ],
-        extra_compile_args=["-std=c++17"],
-    )
-    cmdclass = {"build_ext": CustomBuildExtLinux}
+    cmdclass = {"build_ext": CMakeExtension}
 
 elif system_platform == "Darwin":
-    variables = sysconfig.get_config_vars()
-    # Here we need to switch the deault to MacOs dynamic lib
-    variables["LDSHARED"] = variables["LDSHARED"].replace("-bundle", "-dynamiclib")
-    if sysconfig.get_config_var("LDCXXSHARED"):
-        variables["LDCXXSHARED"] = variables["LDCXXSHARED"].replace("-bundle", "-dynamiclib")
-    custom_calls_extension = Extension(
-        "catalyst.utils.libcustom_calls",
-        sources=[
-            "frontend/catalyst/utils/libcustom_calls.cpp",
-            "frontend/catalyst/utils/jax_cpu_lapack_kernels/lapack_kernels.cpp",
-            "frontend/catalyst/utils/jax_cpu_lapack_kernels/lapack_kernels_using_lapack.cpp",
-        ],
-        extra_compile_args=["-std=c++17"],
-    )
-    cmdclass = {"build_ext": CustomBuildExtMacos}
+    cmdclass = {"build_ext": CMakeExtension}
 
+    # variables = sysconfig.get_config_vars()
+    # Here we need to switch the deault to MacOs dynamic lib
+    # variables["LDSHARED"] = variables["LDSHARED"].replace("-bundle", "-dynamiclib")
+    # if sysconfig.get_config_var("LDCXXSHARED"):
+    #    variables["LDCXXSHARED"] = variables["LDCXXSHARED"].replace("-bundle", "-dynamiclib")
+    # cmdclass = {"build_ext": CustomBuildExtMacos}
 
 project_root_dir = os.path.abspath(os.path.dirname(__file__))
 frontend_dir = os.path.join(project_root_dir, "frontend")
 
 ext_modules = [
-    custom_calls_extension,
+    CMakeExtension("catalyst.utils.libcustom_calls", sourcedir=frontend_dir),
     CMakeExtension("catalyst.utils.wrapper", sourcedir=frontend_dir),
 ]
 
@@ -332,7 +268,7 @@ def run(self):
     package_dir={"": "frontend"},
     include_package_data=True,
     ext_modules=ext_modules,
-    cmdclass=cmdclass,
+    cmdclass={"build_ext": CMakeBuild},
     **description,
     options=options,
 )