Define a Generator framework in Python

cmake/HalideGeneratorHelpers.cmake

@@ -13,6 +13,10 @@ 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_PYTHON_GENERATOR_SOURCE
+                BRIEF_DOCS "Used to store the source file(s) for a Python Generator"
+                FULL_DOCS "Used to store the source file(s) for a Python Generator")
+
 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.")
@@ -68,26 +72,47 @@ function(add_halide_generator TARGET)
             find_package(Halide REQUIRED)
         endif ()
 
-        add_executable(${TARGET} ${ARG_SOURCES})
-        add_executable(${gen} ALIAS ${TARGET})
-        target_link_libraries(${TARGET} PRIVATE Halide::Generator ${ARG_LINK_LIBRARIES})
-
-        if (NOT ARG_NO_DEFAULT_FLAGS AND NOT Halide_NO_DEFAULT_FLAGS)
-            # For crosscompiling builds, the Halide headers will be included using -isystem,
-            # which will cause all warnings to be ignored. This is not helpful, since
-            # we *want* deprecation warnings to be propagated. So we must set
-            # NO_SYSTEM_FROM_IMPORTED in order for it to be seen.
-            set_target_properties(${TARGET} PROPERTIES NO_SYSTEM_FROM_IMPORTED YES)
-            target_compile_options(${TARGET} PRIVATE
-                $<$<CXX_COMPILER_ID:GNU,Clang,AppleClang>:-Wdeprecated-declarations>
-                $<$<CXX_COMPILER_ID:MSVC>:/w14996>  # 4996: compiler encountered deprecated declaration
-            )
-        endif ()
+        if (ARG_SOURCES MATCHES ".py$")
+            if (ARG_LINK_LIBRARIES)
+                message(FATAL_ERROR "You cannot specify LINK_LIBRARIES in conjunction with Python source code.")
+            endif ()
+            if (ARG_PYSTUB)
+                message(FATAL_ERROR "You cannot specify PYSTUB in conjunction with Python source code (only C++ Generators can produce PyStubs).")
+            endif ()
+
+            list(LENGTH ARG_SOURCES len)
+            if (NOT len EQUAL 1)
+                message(FATAL_ERROR "Python Generators must specify exactly one source file.")
+            endif()
+
+            # Make a fake target here that we can attach the Python source to,
+            # so that we can extract 'em in add_halide_library()
+            add_custom_target(${TARGET} ALL)
+            set_property(TARGET ${TARGET} PROPERTY Halide_PYTHON_GENERATOR_SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/${ARG_SOURCES}")
+
+            # TODO: what do we need to do for PACKAGE_NAME PACKAGE_NAMESPACE EXPORT_FILE in this case?
+        else ()
+            add_executable(${TARGET} ${ARG_SOURCES})
+            add_executable(${gen} ALIAS ${TARGET})
+            target_link_libraries(${TARGET} PRIVATE Halide::Generator ${ARG_LINK_LIBRARIES})
+
+            if (NOT ARG_NO_DEFAULT_FLAGS AND NOT Halide_NO_DEFAULT_FLAGS)
+                # For crosscompiling builds, the Halide headers will be included using -isystem,
+                # which will cause all warnings to be ignored. This is not helpful, since
+                # we *want* deprecation warnings to be propagated. So we must set
+                # NO_SYSTEM_FROM_IMPORTED in order for it to be seen.
+                set_target_properties(${TARGET} PROPERTIES NO_SYSTEM_FROM_IMPORTED YES)
+                target_compile_options(${TARGET} PRIVATE
+                    $<$<CXX_COMPILER_ID:GNU,Clang,AppleClang>:-Wdeprecated-declarations>
+                    $<$<CXX_COMPILER_ID:MSVC>:/w14996>  # 4996: compiler encountered deprecated declaration
+                )
+            endif ()
 
-        add_dependencies("${ARG_PACKAGE_NAME}" ${TARGET})
-        export(TARGETS ${TARGET}
-               NAMESPACE ${ARG_PACKAGE_NAMESPACE}
-               APPEND FILE "${ARG_EXPORT_FILE}")
+            add_dependencies("${ARG_PACKAGE_NAME}" ${TARGET})
+            export(TARGETS ${TARGET}
+                   NAMESPACE ${ARG_PACKAGE_NAMESPACE}
+                   APPEND FILE "${ARG_EXPORT_FILE}")
+        endif ()
     endif ()
 
     if (ARG_PYSTUB)
@@ -201,17 +226,32 @@ function(add_halide_library TARGET)
         message(FATAL_ERROR "Missing FROM argument specifying a Halide generator target")
     endif ()
 
-    if (NOT TARGET ${ARG_FROM})
-        # FROM is usually an unqualified name; if we are crosscompiling, we might need a
-        # fully-qualified name, so add the default package name and retry
-        set(FQ_ARG_FROM "${PROJECT_NAME}::halide_generators::${ARG_FROM}")
-        if (NOT TARGET ${FQ_ARG_FROM})
-            message(FATAL_ERROR "Unable to locate FROM as either ${ARG_FROM} or ${FQ_ARG_FROM}")
+    get_property(py_src TARGET ${ARG_FROM} PROPERTY Halide_PYTHON_GENERATOR_SOURCE)
+    if (py_src)
+        # TODO: Python Generators need work to support crosscompiling (https://github.com/halide/Halide/issues/7014)
+        if (NOT TARGET Halide::Python)
+            message(FATAL_ERROR "This version of Halide was built without support for Python bindings; rebuild using WITH_PYTHON_BINDINGS=ON to use this rule with Python Generators.")
         endif ()
-        set(ARG_FROM "${FQ_ARG_FROM}")
-    endif ()
-
-    _Halide_place_dll(${ARG_FROM})
+        if (NOT TARGET Python3::Interpreter)
+            message(FATAL_ERROR "You must call find_package(Python3) in your CMake code in order to use this rule with Python Generators.")
+        endif ()
+        set(PYTHONPATH "$<TARGET_FILE_DIR:Halide::Python>/..")
+        set(GENERATOR_CMD ${CMAKE_COMMAND} -E env PYTHONPATH=${PYTHONPATH} ${Python3_EXECUTABLE} $<SHELL_PATH:${py_src}>)
+        set(GENERATOR_CMD_DEPS ${ARG_FROM} Halide::Python ${py_src})
+    else()
+        if (NOT TARGET ${ARG_FROM})
+            # FROM is usually an unqualified name; if we are crosscompiling, we might need a
+            # fully-qualified name, so add the default package name and retry
+            set(FQ_ARG_FROM "${PROJECT_NAME}::halide_generators::${ARG_FROM}")
+            if (NOT TARGET ${FQ_ARG_FROM})
+                message(FATAL_ERROR "Unable to locate FROM as either ${ARG_FROM} or ${FQ_ARG_FROM}")
+            endif ()
+            set(ARG_FROM "${FQ_ARG_FROM}")
+        endif()
+        set(GENERATOR_CMD "${ARG_FROM}")
+        set(GENERATOR_CMD_DEPS ${ARG_FROM})
+        _Halide_place_dll(${ARG_FROM})
+    endif()
 
     if (ARG_C_BACKEND)
         if (ARG_USE_RUNTIME)
@@ -377,7 +417,7 @@ function(add_halide_library TARGET)
     endif ()
 
     add_custom_command(OUTPUT ${generator_output_files}
-                       COMMAND ${ARG_FROM}
+                       COMMAND ${GENERATOR_CMD}
                        -n "${TARGET}"
                        -d "${gradient_descent}"
                        -g "${ARG_GENERATOR}"
@@ -387,7 +427,7 @@ function(add_halide_library TARGET)
                        -o .
                        "target=$<JOIN:${ARG_TARGETS},$<COMMA>>"
                        ${ARG_PARAMS}
-                       DEPENDS "${ARG_FROM}" ${ARG_PLUGINS}
+                       DEPENDS ${GENERATOR_CMD_DEPS} ${ARG_PLUGINS}
                        VERBATIM)
 
     list(TRANSFORM generator_output_files PREPEND "${CMAKE_CURRENT_BINARY_DIR}/")

python_bindings/CMakeLists.txt

@@ -74,7 +74,7 @@ if (Halide_ASAN_ENABLED)
 endif ()
 
 function(add_python_test)
-    cmake_parse_arguments(ARG "" "FILE;LABEL" "PYTHONPATH;ENVIRONMENT" ${ARGN})
+    cmake_parse_arguments(ARG "" "FILE;LABEL" "PYTHONPATH;ENVIRONMENT;TEST_ARGS" ${ARGN})
 
     list(PREPEND ARG_PYTHONPATH "$<TARGET_FILE_DIR:Halide::Python>/..")
     list(TRANSFORM ARG_PYTHONPATH PREPEND "PYTHONPATH=path_list_prepend:")
@@ -93,7 +93,7 @@ function(add_python_test)
 
     add_test(
         NAME "${test_name}"
-        COMMAND Python3::Interpreter "$<SHELL_PATH:${CMAKE_CURRENT_SOURCE_DIR}/${ARG_FILE}>"
+        COMMAND Python3::Interpreter "$<SHELL_PATH:${CMAKE_CURRENT_SOURCE_DIR}/${ARG_FILE}>" ${ARG_TEST_ARGS}
     )
     set_tests_properties(
         "${test_name}"

python_bindings/src/halide/CMakeLists.txt

@@ -11,6 +11,7 @@ set(native_sources
     PyExternFuncArgument.cpp
     PyFunc.cpp
     PyFuncRef.cpp
+    PyGenerator.cpp
     PyHalide.cpp
     PyImageParam.cpp
     PyInlineReductions.cpp
@@ -33,6 +34,7 @@ list(TRANSFORM native_sources PREPEND "halide_/")
 
 set(python_sources
     __init__.py
+    _generator_helpers.py
     )
 
 # It is technically still possible for a user to override the LIBRARY_OUTPUT_DIRECTORY by setting

python_bindings/src/halide/__init__.py

@@ -1,2 +1,6 @@
 from .halide_ import *
 from .halide_ import _, _1, _2, _3, _4, _5, _6, _7, _8, _9
+from ._generator_helpers import GeneratorParam, InputBuffer, InputScalar, OutputBuffer, \
+     OutputScalar, Generator, alias, generator, active_generator_context, \
+     _get_python_generator_names, _create_python_generator, \
+     _generatorcontext_enter, _generatorcontext_exit