Remove add_python_aot_extension() rule in CMake

README_cmake.md

@@ -807,6 +807,7 @@ add_halide_library(<target> FROM <generator-target>
                    [PLUGINS plugin1 [plugin2 ...]]
                    [AUTOSCHEDULER scheduler-name]
                    [GRADIENT_DESCENT]
+                   [PYTHON_EXTENSION_LIBRARY]
                    [C_BACKEND]
                    [REGISTRATION OUTVAR]
                    [HEADER OUTVAR]
@@ -869,6 +870,13 @@ 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

README_python.md

@@ -172,7 +172,7 @@ objects without any explicit conversion necessary.
 
 ## Using Halide Generators from Python
 
-### Compiling a C++ Generator for Python
+### Using a C++ Generator from Python
 
 Let's say we have a very simple Generator:
 
@@ -192,12 +192,13 @@ HALIDE_REGISTER_GENERATOR(MyFilter, my_filter)
 ```
 
 If you are using CMake, the simplest thing is to use
-`add_python_aot_extension()`, defined in PythonExtensionHelpers.cmake:
+`add_halide_library()` (defined in HalideGeneratorHelpers.cmake) with the `PYTHON_EXTENSION_LIBRARY` option:
 
 ```
-add_python_aot_extension(my_filter
+add_halide_library(my_filter
                          GENERATOR my_filter_generator
                          SOURCES my_filter_generator.cpp
+                         PYTHON_EXTENSION_LIBRARY
                          [ FEATURES ... ]
                          [ PARAMS ... ])
 ```

cmake/HalideGeneratorHelpers.cmake

@@ -1,6 +1,7 @@
 cmake_minimum_required(VERSION 3.22)
 
 include(${CMAKE_CURRENT_LIST_DIR}/HalideTargetHelpers.cmake)
+include(${CMAKE_CURRENT_LIST_DIR}/TargetExportScript.cmake)
 
 define_property(TARGET PROPERTY Halide_RT_TARGETS
                 BRIEF_DOCS "On a Halide runtime target, lists the targets the runtime backs"
@@ -136,7 +137,7 @@ function(add_halide_library TARGET)
     # Parse the arguments and set defaults for missing values.
     ##
 
-    set(options C_BACKEND GRADIENT_DESCENT)
+    set(options C_BACKEND GRADIENT_DESCENT PYTHON_EXTENSION_LIBRARY)
     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})
@@ -268,6 +269,14 @@ 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})
@@ -346,6 +355,14 @@ 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}")
+    endif ()
+
 endfunction()
 
 ##
@@ -511,3 +528,17 @@ function(_Halide_fix_xcode TARGET)
         target_sources("${TARGET}" PRIVATE "${empty_file}")
     endif ()
 endfunction()
+
+function(_Halide_target_export_single_symbol TARGET SYMBOL)
+    file(WRITE
+         "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript.apple"
+         "_${SYMBOL}\n")
+    file(WRITE
+         "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript"
+         "{ global: ${SYMBOL}; local: *; };\n")
+    target_export_script(
+        ${TARGET}
+        APPLE_LD "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript.apple"
+        GNU_LD "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript"
+    )
+endfunction()

cmake/HalidePythonExtensionHelpers.cmake

@@ -5,30 +5,6 @@ include(${CMAKE_CURRENT_LIST_DIR}/TargetExportScript.cmake)
 
 set(_STUB_DIR "${Halide_SOURCE_DIR}/python_bindings/stub")
 
-# There are two sorts of Python Extensions that we can produce for a Halide Generator
-# written in C++:
-#
-# - One that is essentially the 'native code' output of a Generator, wrapped with enough CPython
-#   glue code to make it callable from Python. This is analogous to the usual Generator output
-#   when building a C++ codebase, and is the usual mode used for distribution of final product;
-#   these correspond to 'ahead-of-time' (AOT) code generation. The resulting code has no dependency
-#   on libHalide. We'll refer to this sort of extension as an "AOT extension".
-#
-# - One that essentially *the Generator itself*, wrapped in CPython glue code to make it callable
-#   from Python at Halide compilation time. This is analogous to the (rarely used) GeneratorStub
-#   code that can be used to compose multiple Generators together. The resulting extension *does*
-#   depend on libHalide, and can be used in either JIT or AOT mode for compilation.
-#   We'll refer to this sort of extension as a "Stub extension".
-#
-# For testing purposes here, we don't bother using distutils/setuptools to produce a properly-packaged
-# Python extension; rather, we simply produce a .so file with the correct name exported, and ensure
-# it's in the PYTHONPATH when testing.
-#
-# In our build files here, we build both kinds of extension for every Generator in the generators/
-# directory (even though not all are used). As a simplistic way to distinguish between the two
-# sorts of extensions, we use the unadorned Generator name for AOT extensions, and the Generator name
-# suffixed with "_stub" for Stub extensions. (TODO: this is unsatisfyingly hackish; better suggestions
-# would be welcome.)
 
 function(_target_export_single_symbol TARGET SYMBOL)
     file(WRITE
@@ -44,47 +20,6 @@ function(_target_export_single_symbol TARGET SYMBOL)
     )
 endfunction()
 
-function(add_python_aot_extension TARGET)
-    set(options)
-    set(oneValueArgs GENERATOR FUNCTION_NAME)
-    set(multiValueArgs SOURCES LINK_LIBRARIES FEATURES PARAMS)
-    cmake_parse_arguments(ARG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
-
-    if (NOT ARG_GENERATOR)
-        set(ARG_GENERATOR "${TARGET}")
-    endif ()
-
-    if (NOT ARG_FUNCTION_NAME)
-        set(ARG_FUNCTION_NAME "${ARG_GENERATOR}")
-    endif ()
-
-    # Create the Halide generator executable.
-    add_executable(${TARGET}.generator ${ARG_SOURCES})
-    target_link_libraries(${TARGET}.generator PRIVATE Halide::Generator ${ARG_LINK_LIBRARIES})
-
-    # TODO: this should work (and would be preferred to the code above)
-    # but CMake fails with "targets not yet defined"; investigate.
-    # add_halide_generator(${TARGET}.generator
-    #                      SOURCES ${ARG_SOURCES})
-
-    # Run the Generator to produce a static library of AOT code,
-    # plus the 'python_extension' code necessary to produce a useful
-    # AOT Extention for Python:
-    add_halide_library(aot_${TARGET}
-                       FROM ${TARGET}.generator
-                       GENERATOR ${ARG_GENERATOR}
-                       FUNCTION_NAME ${ARG_FUNCTION_NAME}
-                       PYTHON_EXTENSION ${TARGET}.py.cpp
-                       FEATURES ${ARG_FEATURES}
-                       PARAMS ${ARG_PARAMS}
-                       TARGETS cmake)
-
-    Python3_add_library(${TARGET} MODULE WITH_SOABI ${${TARGET}.py.cpp})
-    target_link_libraries(${TARGET} PRIVATE aot_${TARGET})
-    set_target_properties(${TARGET} PROPERTIES OUTPUT_NAME ${ARG_FUNCTION_NAME})
-    _target_export_single_symbol(${TARGET} "PyInit_${ARG_FUNCTION_NAME}")
-endfunction()
-
 function(add_python_stub_extension TARGET)
     set(options)
     set(oneValueArgs GENERATOR MODULE)

python_bindings/test/CMakeLists.txt

@@ -1,3 +1,8 @@
+if (TARGET_WEBASSEMBLY AND Halide_TARGET MATCHES "wasm")
+    message(WARNING "Python tests are skipped under WASM.")
+    return()
+endif ()
+
 add_subdirectory(apps)
 add_subdirectory(correctness)
 add_subdirectory(generators)

python_bindings/test/generators/CMakeLists.txt

@@ -27,11 +27,19 @@ set(GENPARAMS_simple
     func_input.type=uint8)
 
 foreach (GEN IN LISTS GENERATORS)
-    add_python_aot_extension(py_aot_${GEN}
-                             GENERATOR ${GEN}
-                             FEATURES ${FEATURES_${GEN}}
-                             PARAMS ${GENPARAMS_${GEN}}
-                             SOURCES ${GEN}_generator.cpp)
+    add_halide_generator(py_aot_${GEN}.generator
+                         SOURCES ${GEN}_generator.cpp)
+
+    add_halide_library(py_aot_${GEN}
+                       FROM py_aot_${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}}
+                       PARAMS ${GENPARAMS_${GEN}})
 
     add_python_stub_extension(py_stub_${GEN}
                               SOURCES ${GEN}_generator.cpp