Add add_halide_python_extension_library() rule (halide#6979)

* Add add_halide_python_extension_library() rule

This adds a rule to create a single Python extension library from one (or more) halide_library rules. This allows you to package multiple Halide filters into a single Python module, which is nice because (1) being able to organize is good, and (2) all the filters in a single Python extension module share the same Halide runtime, including (e.g.) thread pools and method overrides.

(It also removes the just-recently-added PYTHON_EXTENSION_LIBRARY option from the add_halide_library rule, as this new rule is better and more flexible in pretty much every way.)

This modifies the content of our `python_extension` output in such a way that existing uses should be completely unaffected, but defining the right preprocessor macros allows us to split the function wrappers up from the method-definition declaration, so we don't have to generate any new code artifiacts to make this work.

Partially addresses halide#6956.

* Omits -D in target_compile_definitions

* be explicit about setting to empty

* Add quotes

* Add comments re BUILD_INTERFACE

* Add MODULE_NAME comment

* Remove "defined in HalideGeneratorHelpers.cmake"

* Add comment re add_halide_runtime()

* osx, macos, darwin, oh m

* blankity blank blank

* Use OBJECT library instead

* Add comment about X-macros

* Update HalideGeneratorHelpers.cmake

README_cmake.md

@@ -807,7 +807,6 @@ add_halide_library(<target> FROM <generator-target>
                    [PLUGINS plugin1 [plugin2 ...]]
                    [AUTOSCHEDULER scheduler-name]
                    [GRADIENT_DESCENT]
-                   [PYTHON_EXTENSION_LIBRARY]
                    [C_BACKEND]
                    [REGISTRATION OUTVAR]
                    [HEADER OUTVAR]
@@ -870,13 +869,6 @@ gradient descent calculation in TensorFlow or PyTorch. See
 `Generator::build_gradient_module()` for more documentation. This corresponds to
 passing `-d 1` at the generator command line.
 
-If `PYTHON_EXTENSION_LIBRARY` is set, then a Python Extension will be built that
-wraps the C/C++ call with CPython glue to allow use of the generated code from
-Python 3.x. The result will be a a shared library of the form
-`<target>.<soabi>.so`, where <soabi> describes the specific Python version and
-platform (e.g., `cpython-310-darwin` for Python 3.10 on OSX.) See
-`README_python.md` for examples of use.
-
 If the `C_BACKEND` option is set, this command will invoke the configured C++
 compiler on a generated source. Note that a `<target>.runtime` target is _not_
 created in this case, and the `USE_RUNTIME` option is ignored. Other options
@@ -945,9 +937,33 @@ and [apps/hannk](https://github.com/halide/Halide/tree/master/apps/hannk) for a
 If `PYSTUB` is specified, then a Python Extension will be built that
 wraps the Generator with CPython glue to allow use of the Generator
 Python 3.x. The result will be a a shared library of the form
-`<target>_pystub.<soabi>.so`, where <soabi> describes the specific Python version and platform (e.g., `cpython-310-darwin` for Python 3.10 on OSX.) See
+`<target>_pystub.<soabi>.so`, where <soabi> describes the specific Python version and platform (e.g., `cpython-310-darwin` for Python 3.10 on macOS.) See
 `README_python.md` for examples of use.
 
+#### `add_halide_python_extension_library`
+
+This function wraps the outputs of one or more `add_halide_library` targets with glue code to produce
+a Python Extension library.
+
+```
+add_halide_python_extension_library(
+    target
+    [MODULE_NAME module-name]
+    HALIDE_LIBRARIES library1 ...
+)
+```
+
+The `MODULE_NAME` argument specifies the name of the Python module that will be created. If omitted,
+it defaults to `target`.
+
+`HALIDE_LIBRARIES` is a list of one of more `add_halide_library` targets. Each will be added to the
+extension as a callable method of the module. Note that every library specified must be built with
+the `PYTHON_EXTENSION` keyword specified, and all libraries must use the same Halide runtime.
+
+The result will be a a shared library of the form
+`<target>.<soabi>.so`, where <soabi> describes the specific Python version and
+platform (e.g., `cpython-310-darwin` for Python 3.10 on macOS.)
+
 ## Cross compiling
 
 Cross-compiling in CMake can be tricky, since CMake doesn't easily support

README_python.md

@@ -192,29 +192,52 @@ HALIDE_REGISTER_GENERATOR(MyFilter, my_filter)
 ```
 
 If you are using CMake, the simplest thing is to use
-`add_halide_library()` (defined in HalideGeneratorHelpers.cmake) with the `PYTHON_EXTENSION_LIBRARY` option:
+`add_halide_library` and `add_halide_python_extension_library()`:
 
 ```
+# Build a Halide library as you usually would, but be sure to include `PYTHON_EXTENSION`
 add_halide_library(my_filter
-                         GENERATOR my_filter_generator
-                         SOURCES my_filter_generator.cpp
-                         PYTHON_EXTENSION_LIBRARY
-                         [ FEATURES ... ]
-                         [ PARAMS ... ])
+                   GENERATOR my_filter_generator
+                   SOURCES my_filter_generator.cpp
+                   PYTHON_EXTENSION output_path_var
+                   [ FEATURES ... ]
+                   [ PARAMS ... ])
+
+# Now wrap the generated code with a Python extension.
+# (Note that module name defaults to match the target name; we only
+# need to specify MODULE_NAME if we need a name that may differ)
+add_halide_python_extension_library(my_extension
+                                    MODULE_NAME my_module
+                                    HALIDE_LIBRARIES my_filter)
 ```
 
 This compiles the Generator code in `my_filter_generator.cpp` with the
-registered name `my_filter` to produce the target `my_filter`, which is a Python
-Extension in the form of a shared library (e.g.,
-`foo.cpython-310-x86_64-linux-gnu.so`).
+registered name `my_filter` to produce the target `my_filter`, and then wraps
+the compiled output with a Python extension. The result will be a shared library of the form
+`<target>.<soabi>.so`, where <soabi> describes the specific Python version and
+platform (e.g., `cpython-310-darwin` for Python 3.10 on OSX.)
+
+Note that you can combine multiple Halide libraries into a single Python module;
+this is convenient for packagaing, but also because all the libraries in a single
+extension module share the same Halide runtime (and thus, the same caches, thread pools, etc.):
+
+```
+add_halide_library(my_filter1 ...)
+add_halide_library(my_filter2 ...)
+add_halide_library(my_filter3 ...)
+
+add_halide_python_extension_library(my_extension
+                                    MODULE_NAME my_module
+                                    HALIDE_LIBRARIES my_filter my_filter2 my_filter3)
+```
 
 ### Calling a C++ Generator from Python
 
 As long as the shared library is in `PYTHONPATH`, it can be imported and used
 directly. For the example above:
 
 ```
-from my_filter import my_filter
+from my_module import my_filter
 import imageio
 import numpy as np
 

cmake/HalideGeneratorHelpers.cmake

@@ -13,6 +13,18 @@ define_property(TARGET PROPERTY Halide_GENERATOR_HAS_POST_BUILD
                 BRIEF_DOCS "On a Halide generator target, true if Halide.dll copy command has already been added."
                 FULL_DOCS "On a Halide generator target, true if Halide.dll copy command has already been added.")
 
+define_property(TARGET PROPERTY Halide_LIBRARY_RUNTIME_TARGET
+                BRIEF_DOCS "On a Halide library target, the runtime it uses."
+                FULL_DOCS "On a Halide library target, the runtime it uses.")
+
+define_property(TARGET PROPERTY Halide_LIBRARY_PYTHON_EXTENSION_CPP
+                BRIEF_DOCS "On a Halide library target, the .py.cpp generated for it (absent if none)."
+                FULL_DOCS "On a Halide library target, the .py.cpp generated for it (absent if none).")
+
+define_property(TARGET PROPERTY Halide_LIBRARY_FUNCTION_NAME
+                BRIEF_DOCS "On a Halide library target, the FUNCTION_NAME used."
+                FULL_DOCS "On a Halide library target, the FUNCTION_NAME used.")
+
 ##
 # Function to simplify writing the CMake rules for creating a generator executable
 # that follows our recommended cross-compiling workflow.
@@ -175,7 +187,7 @@ function(add_halide_library TARGET)
     # Parse the arguments and set defaults for missing values.
     ##
 
-    set(options C_BACKEND GRADIENT_DESCENT PYTHON_EXTENSION_LIBRARY)
+    set(options C_BACKEND GRADIENT_DESCENT)
     set(oneValueArgs FROM GENERATOR FUNCTION_NAME NAMESPACE USE_RUNTIME AUTOSCHEDULER HEADER ${extra_output_names})
     set(multiValueArgs TARGETS FEATURES PARAMS PLUGINS)
     cmake_parse_arguments(ARG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
@@ -307,14 +319,6 @@ function(add_halide_library TARGET)
     endif ()
     list(APPEND generator_output_files ${generator_sources})
 
-    # If we're building an extension library, we also need the extension .cpp
-    # file, so quietly add it if it wasn't already specified
-    if (ARG_PYTHON_EXTENSION_LIBRARY)
-        if (NOT ARG_PYTHON_EXTENSION)
-            set(ARG_PYTHON_EXTENSION "${TARGET}.py.cpp")
-        endif()
-    endif ()
-
     # Add in extra outputs using the table defined at the start of this function
     foreach (out IN LISTS extra_output_names)
         if (ARG_${out})
@@ -394,13 +398,85 @@ function(add_halide_library TARGET)
     target_include_directories("${TARGET}" INTERFACE "$<BUILD_INTERFACE:${CMAKE_CURRENT_BINARY_DIR}>")
     target_link_libraries("${TARGET}" INTERFACE "${ARG_USE_RUNTIME}")
 
-    if (ARG_PYTHON_EXTENSION_LIBRARY)
-        Python3_add_library(${TARGET}_py_ext_lib MODULE WITH_SOABI ${ARG_PYTHON_EXTENSION})
-        target_link_libraries(${TARGET}_py_ext_lib PRIVATE ${TARGET})
-        set_target_properties(${TARGET}_py_ext_lib PROPERTIES OUTPUT_NAME ${ARG_FUNCTION_NAME})
-        _Halide_target_export_single_symbol(${TARGET}_py_ext_lib "PyInit_${ARG_FUNCTION_NAME}")
+    # Save some info for add_halide_python_extension_library() in case it is used for this target.
+    set_property(TARGET "${TARGET}" PROPERTY Halide_LIBRARY_RUNTIME_TARGET "${ARG_USE_RUNTIME}")
+    set_property(TARGET "${TARGET}" PROPERTY Halide_LIBRARY_FUNCTION_NAME "${ARG_FUNCTION_NAME}")
+    if ("python_extension" IN_LIST generator_outputs)
+        set_property(TARGET "${TARGET}" PROPERTY Halide_LIBRARY_PYTHON_EXTENSION_CPP "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.py.cpp")
+    endif ()
+endfunction()
+
+function(add_halide_python_extension_library TARGET)
+    set(options "")
+    set(oneValueArgs MODULE_NAME)
+    set(multiValueArgs HALIDE_LIBRARIES)
+    cmake_parse_arguments(ARG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
+
+    if (NOT ARG_MODULE_NAME)
+        set(ARG_MODULE_NAME "${TARGET}")
     endif ()
 
+    if (NOT ARG_HALIDE_LIBRARIES)
+        message(FATAL_ERROR "HALIDE_LIBRARIES must be specified")
+    endif ()
+
+    set(runtimes "")
+    set(pycpps "")
+    set(function_names "")  # space-separated X-macros
+    foreach (lib IN LISTS ARG_HALIDE_LIBRARIES)
+        if (NOT TARGET "${lib}")
+            message(FATAL_ERROR "${lib} is not a valid target")
+        endif ()
+
+        get_property(runtime_used TARGET ${lib} PROPERTY Halide_LIBRARY_RUNTIME_TARGET)
+        if (NOT runtime_used)
+            message(FATAL_ERROR "${lib} does not appear to have a Halide Runtime specified")
+        endif ()
+        list(APPEND runtimes ${runtime_used})
+
+        get_property(function_name TARGET ${lib} PROPERTY Halide_LIBRARY_FUNCTION_NAME)
+        if (NOT function_name)
+            message(FATAL_ERROR "${lib} does not appear to have a Function name specified")
+        endif ()
+        # Strip C++ namespace(s), if any
+        string(REGEX REPLACE ".*::(.*)" "\\1" function_name "${function_name}")
+        string(APPEND function_names " X(${function_name})")
+
+        get_property(pycpp TARGET ${lib} PROPERTY Halide_LIBRARY_PYTHON_EXTENSION_CPP)
+        if (NOT pycpp)
+            message(FATAL_ERROR "${lib} must be built with PYTHON_EXTENSION specified in order to use it with add_halide_python_extension_library()")
+        endif ()
+        list(APPEND pycpps ${pycpp})
+    endforeach ()
+
+    list(REMOVE_DUPLICATES runtimes)
+    list(LENGTH runtimes len)
+    if (NOT len EQUAL 1)
+        message(FATAL_ERROR "${TARGET} requires all libraries to use the same Halide Runtime, but saw ${len}: ${runtimes}")
+    endif ()
+
+    # Module def code is the same in all of them, so arbitrarily choose the first
+    list(GET pycpps 0 module_definition_cpp)
+    add_library(${TARGET}_module_definition OBJECT ${module_definition_cpp})
+    set_target_properties(${TARGET}_module_definition PROPERTIES
+                          CXX_VISIBILITY_PRESET hidden
+                          VISIBILITY_INLINES_HIDDEN ON
+                          POSITION_INDEPENDENT_CODE ON)
+    target_link_libraries(${TARGET}_module_definition PRIVATE Halide::Runtime Python3::Module)
+    # Compile it with the right preprocessor definitions to provide the module defs,
+    # but not the function implementations
+    target_compile_definitions(${TARGET}_module_definition PRIVATE
+                               HALIDE_PYTHON_EXTENSION_OMIT_FUNCTION_DEFINITIONS
+                               HALIDE_PYTHON_EXTENSION_MODULE=${ARG_MODULE_NAME}
+                               "HALIDE_PYTHON_EXTENSION_FUNCTIONS=${function_names}")
+
+    # Now compile all the pycpps to build the function implementations (but not the module def)
+    Python3_add_library(${TARGET} MODULE WITH_SOABI ${pycpps} $<TARGET_OBJECTS:${TARGET}_module_definition>)
+    target_link_libraries(${TARGET} PRIVATE ${ARG_HALIDE_LIBRARIES})
+    target_compile_definitions(${TARGET} PRIVATE
+                               HALIDE_PYTHON_EXTENSION_OMIT_MODULE_DEFINITION)
+    set_target_properties(${TARGET} PROPERTIES OUTPUT_NAME ${ARG_MODULE_NAME})
+    _Halide_target_export_single_symbol(${TARGET} "PyInit_${ARG_MODULE_NAME}")
 endfunction()
 
 ##

python_bindings/test/correctness/CMakeLists.txt

@@ -15,6 +15,7 @@ set(tests
     extern.py
     float_precision_test.py
     iroperator.py
+    multi_method_module_test.py
     multipass_constraints.py
     pystub.py
     rdom.py

python_bindings/test/correctness/multi_method_module_test.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+import multi_method_module
+
+def test_simple():
+    buffer_input = np.ndarray([2, 2], dtype=np.uint8)
+    buffer_input[0, 0] = 123
+    buffer_input[0, 1] = 123
+    buffer_input[1, 0] = 123
+    buffer_input[1, 1] = 123
+
+    func_input = np.ndarray([2, 2], dtype=np.uint8)
+    func_input[0, 0] = 0
+    func_input[0, 1] = 1
+    func_input[1, 0] = 1
+    func_input[1, 1] = 2
+
+    float_arg = 3.5
+
+    simple_output = np.ndarray([2, 2], dtype=np.float32)
+
+    multi_method_module.simple(buffer_input, func_input, float_arg, simple_output)
+
+    assert simple_output[0, 0] == 3.5 + 123
+    assert simple_output[0, 1] == 4.5 + 123
+    assert simple_output[1, 0] == 4.5 + 123
+    assert simple_output[1, 1] == 5.5 + 123
+
+def test_user_context():
+    output = bytearray("\0\0\0\0", "ascii")
+    multi_method_module.user_context(None, ord('q'), output)
+    assert output == bytearray("qqqq", "ascii")
+
+
+if __name__ == "__main__":
+    test_simple()
+    test_user_context()

python_bindings/test/generators/CMakeLists.txt

@@ -24,6 +24,15 @@ set(GENPARAMS_complex
 set(GENPARAMS_simple
     func_input.type=uint8)
 
+# Since simple and user_context are going to be bound into a single
+# Python extension library, they must share the same Halide runtime.
+# Since none of these use features that require custom runtimes,
+# we'll just have user_context use the runtime we build for simple.
+#
+# TODO: replace this hackery with appropriate use of add_halide_runtime()
+# (See https://github.com/halide/Halide/issues/6981)
+set(RUNTIME_user_context py_aot_simple.runtime)
+
 foreach (GEN IN LISTS GENERATORS)
     add_halide_generator(py_gen_${GEN}.generator
                          SOURCES ${GEN}_generator.cpp
@@ -33,10 +42,19 @@ foreach (GEN IN LISTS GENERATORS)
                        FROM py_gen_${GEN}.generator
                        GENERATOR ${GEN}
                        FUNCTION_NAME ${GEN}
-                       # Note that PYTHON_EXTENSION_LIBRARY doesn't take
-                       # an argument; it will always produce a file of the form
-                       # ${TARGET}.${Python3_SOABI}.so (or .pyd on Windows)
-                       PYTHON_EXTENSION_LIBRARY
-                       FEATURES ${FEATURES_${GEN}}
+                       USE_RUNTIME ${RUNTIME_${GEN}}
+                       PYTHON_EXTENSION _ignored_result
+                       FEATURES ${FEATURES_${GEN}} c_plus_plus_name_mangling
                        PARAMS ${GENPARAMS_${GEN}})
+
+    add_halide_python_extension_library(pyext_${GEN}
+                                        MODULE_NAME ${GEN}
+                                        HALIDE_LIBRARIES py_aot_${GEN})
+
 endforeach ()
+
+# Bind several libraries into a single python extension;
+# they will be in the same module (and share the same Halide runtime)
+add_halide_python_extension_library(pyext_multi_method_module
+                                    MODULE_NAME multi_method_module
+                                    HALIDE_LIBRARIES py_aot_simple py_aot_user_context)