Clean up Python extensions in python_bindings (halide#6670)

* Remove the nobuild/partialbuildmethod tests from python_bindings/

They no longer serve a purpose and are redundant to other tests.

* WIP

* Update pystub.py

* wip

* wip

* wip

* Update TargetExportScript.cmake

* Update PythonExtensionHelpers.cmake

* PyExtensionGen didn't handle zero-dimensional buffers

cmake/PythonExtensionHelpers.cmake

@@ -0,0 +1,123 @@
+include(HalideGeneratorHelpers)
+include(TargetExportScript)
+
+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)
+    configure_file("${_STUB_DIR}/ext.ldscript.apple.in" "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript.apple")
+    configure_file("${_STUB_DIR}/ext.ldscript.linux.in" "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript")
+    target_export_script(
+        ${TARGET}
+        APPLE_LD "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript.apple"
+        GNU_LD "${CMAKE_CURRENT_BINARY_DIR}/${TARGET}.ldscript"
+    )
+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)
+
+    # Take the native-code output of the Generator, add the Python-Extension
+    # code (to make it callable from Python), and build it into the AOT Extension we need.
+    if (CMAKE_VERSION VERSION_GREATER_EQUAL 3.17)
+        # Add soabi info (like cpython-310-x86_64-linux-gnu)
+        # when CMake is new enough to know how to do it.
+        set(abi_flags WITH_SOABI)
+    else ()
+        set(abi_flags "")
+    endif ()
+
+    Python3_add_library(${TARGET} MODULE ${abi_flags} ${${TARGET}.py.cpp})
+    target_link_libraries(${TARGET} PRIVATE aot_${TARGET})
+    set_target_properties(${TARGET} PROPERTIES OUTPUT_NAME ${ARG_GENERATOR})
+    target_export_single_symbol(${TARGET} ${ARG_FUNCTION_NAME})
+endfunction()
+
+function(add_python_stub_extension TARGET)
+    set(options)
+    set(oneValueArgs GENERATOR MODULE)
+    set(multiValueArgs SOURCES LINK_LIBRARIES)
+    cmake_parse_arguments(ARG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
+
+    if (NOT ARG_GENERATOR)
+        set(ARG_GENERATOR "${TARGET}")
+    endif ()
+
+    if (NOT ARG_MODULE)
+        set(ARG_MODULE "${TARGET}_stub")
+    endif ()
+
+    # Produce a Stub Extension for the same Generator:
+    # Compiling PyStub.cpp, then linking with the generator's .o file, PyStubImpl.o,
+    # plus the same libHalide being used by halide.so.
+    #
+    # Note that we set HALIDE_PYSTUB_MODULE_NAME to $*_stub (e.g. foo_stub) but
+    # set HALIDE_PYSTUB_GENERATOR_NAME to the unadorned name of the Generator.
+    Python3_add_library(${TARGET} MODULE ${_STUB_DIR}/PyStub.cpp ${ARG_SOURCES})
+    set_target_properties(${TARGET} PROPERTIES
+                          CXX_VISIBILITY_PRESET hidden
+                          VISIBILITY_INLINES_HIDDEN ON
+                          POSITION_INDEPENDENT_CODE ON)
+    target_compile_definitions(${TARGET} PRIVATE
+                               "HALIDE_PYSTUB_GENERATOR_NAME=${ARG_GENERATOR}"
+                               "HALIDE_PYSTUB_MODULE_NAME=${ARG_MODULE}")
+    target_link_libraries(${TARGET} PRIVATE Halide::PyStubs ${ARG_LINK_LIBRARIES})
+    set_target_properties(${TARGET} PROPERTIES OUTPUT_NAME ${ARG_MODULE})
+    target_export_single_symbol(${TARGET} ${ARG_MODULE})
+endfunction()

cmake/TargetExportScript.cmake

@@ -10,7 +10,7 @@ function(target_export_script TARGET)
     cmake_parse_arguments(ARG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
 
     get_property(target_type TARGET ${TARGET} PROPERTY TYPE)
-    if (NOT target_type STREQUAL "SHARED_LIBRARY")
+    if (NOT target_type STREQUAL "SHARED_LIBRARY" AND NOT target_type STREQUAL "MODULE_LIBRARY")
         # Linker scripts do nothing on non-shared libraries.
         return()
     endif ()

python_bindings/CMakeLists.txt

@@ -19,7 +19,7 @@ endif ()
 ##
 
 add_subdirectory(src)
-include(stub/CMakeLists.txt)
+add_subdirectory(stub)
 
 option(WITH_TEST_PYTHON "Build Python tests" ON)
 if (WITH_TESTS AND WITH_TEST_PYTHON)

python_bindings/Makefile

@@ -68,103 +68,138 @@ $(BIN)/src/%.o: $(ROOT_DIR)/src/%.cpp
 	@mkdir -p $(@D)
 	@$(CXX) $(CCFLAGS) -c $< -o $@
 
-
-$(BIN)/%_generator.o: $(ROOT_DIR)/correctness/%_generator.cpp $(HALIDE_DISTRIB_PATH)/include/Halide.h
-	@echo Building $@...
-	@mkdir -p $(@D)
-	@$(CXX) $(CCFLAGS) -c $< -o $@
-
-$(BIN)/PyStubImpl.o: $(ROOT_DIR)/stub/PyStubImpl.cpp $(HALIDE_DISTRIB_PATH)/include/Halide.h
+$(BIN)/%_generator.o: $(ROOT_DIR)/correctness/generators/%_generator.cpp $(HALIDE_DISTRIB_PATH)/include/Halide.h
 	@echo Building $@...
 	@mkdir -p $(@D)
 	@$(CXX) $(CCFLAGS) -c $< -o $@
 
-# Produce a Python extension for a C++ generator by compiling PyStub.cpp
-# (with HALIDE_PYSTUB_GENERATOR_NAME defined to the Generator's build name),
-# and linking with the generator's .o file, PyStubImpl.o, plus the same libHalide
-# being used by halide.so.
+# There are two sorts of Python Extensions that we can produce for a Halide Generator
+# written in C++:
 #
-# You can optionally also define HALIDE_PYSTUB_MODULE_NAME if you want the Python
-# module name to be something other than the Generator build name.
-$(BIN)/%_PyStub.o: $(ROOT_DIR)/stub/PyStub.cpp
-	@echo Building $@...
-	@mkdir -p $(@D)
-	@$(CXX) $(CCFLAGS) -DHALIDE_PYSTUB_GENERATOR_NAME=$* -c $< -o $@
+# - 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.)
 
-$(BIN)/generators/%.so: $(BIN)/%_PyStub.o $(BIN)/PyStubImpl.o $(BIN)/%_generator.o $(LIBHALIDE)
+$(BIN)/PyStubImpl.o: $(ROOT_DIR)/stub/PyStubImpl.cpp $(HALIDE_DISTRIB_PATH)/include/Halide.h
 	@echo Building $@...
 	@mkdir -p $(@D)
-	@$(CXX) $^ $(LDFLAGS) -shared -o $@
+	@$(CXX) $(CCFLAGS) -c $< -o $@
 
 # Compile the generators:
-$(BIN)/%.gen: $(HALIDE_DISTRIB_PATH)/tools/GenGen.cpp $(BIN)/%_generator.o $(LIBHALIDE)
+$(BIN)/%.generator: $(HALIDE_DISTRIB_PATH)/tools/GenGen.cpp $(BIN)/%_generator.o $(LIBHALIDE)
 	@echo Building $@...
 	@mkdir -p $(@D)
 	@$(CXX) $(CCFLAGS) $(LDFLAGS) $^ -o $@
 
 # Special generator for generating a runtime:
-$(BIN)/runtime.gen: $(HALIDE_DISTRIB_PATH)/tools/GenGen.cpp $(LIBHALIDE)
+$(BIN)/runtime.generator: $(HALIDE_DISTRIB_PATH)/tools/GenGen.cpp $(LIBHALIDE)
 	@echo Building $@...
 	@mkdir -p $(@D)
 	@$(CXX) $(CCFLAGS) $(LDFLAGS) $^ -o $@
 
 # Generate a runtime:
-$(BIN)/runtime.a: $(BIN)/runtime.gen
+$(BIN)/runtime.a: $(BIN)/runtime.generator
 	@echo Building $@...
 	@mkdir -p $(@D)
 	@$< -r runtime -o $(BIN) target=host
 
-# Which target features to use for which test targets.
-target_features_addconstant=-no_runtime
-target_features_bit=-no_runtime
-target_features_user_context=-user_context-no_runtime
+# Construct linker script that will export *just* the PyInit entry
+# point we want. (If we don't do this we can have interesting failures
+# when loading multiple of these Python extensions in the same space.)
+ifeq ($(UNAME), Darwin)
+$(BIN)/%.ldscript: $(ROOT_DIR)/stub/ext.ldscript.apple.in
+	@echo Building $@...
+	@mkdir -p $(@D)
+	@cat $< | sed 's/$${SYMBOL}/$*/' > $@
+PYEXT_LDSCRIPT_FLAG = -Wl,-exported_symbols_list %LDSCRIPT%
+else
+# Assume Desktop Linux
+$(BIN)/%.ldscript: $(ROOT_DIR)/stub/ext.ldscript.linux.in
+	@echo Building $@...
+	@mkdir -p $(@D)
+	@cat $< | sed 's/$${SYMBOL}/$*/' > $@
+PYEXT_LDSCRIPT_FLAG = -Wl,--version-script=%LDSCRIPT%
+endif
 
-# Make the generator generate a Python extension:
-$(BIN)/%.py.cpp $(BIN)/%.a $(BIN)/%.h: $(BIN)/%.gen
+# Some Generators require extra Halide Target Features to be set.
+FEATURES_user_context=-user_context
+
+# Some Generators have undefined types, sizes, etc that are useful for Stubs extensions,
+# but unacceptable for AOT Extensions; ensure that all of those are explicitly
+# specified for AOT. (We currently don't use or test these in AOT form, so the settings
+# are somewhat arbitrary.)
+GENPARAMS_complex=\
+	array_input.size=2 \
+	array_input.type=uint8 \
+	int_arg.size=2 \
+	simple_input.type=uint8 \
+	untyped_buffer_input.type=uint8
+
+GENPARAMS_simple=\
+	func_input.type=uint8
+
+# 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:
+$(BIN)/%.py.cpp $(BIN)/%.a $(BIN)/%.h: $(BIN)/%.generator
 	@echo Building $@...
 	@LD_LIBRARY_PATH=$(HALIDE_DISTRIB_PATH)/bin $< \
 	    -e static_library,c_header,python_extension \
-	    -g $(notdir $(basename $<)) -o $(BIN) \
-	    target=host$(target_features_$(notdir $(basename $<)))
+	    -g $(notdir $(basename $<)) \
+	    -o $(BIN) \
+	    target=host-no_runtime$(FEATURES_$(notdir $(basename $<))) \
+	    $(GENPARAMS_$(notdir $(basename $<)))
 
 # Compile the generated Python extension(s):
 $(BIN)/%.py.o: $(BIN)/%.py.cpp
 	@echo Building $@...
 	@$(CXX) -c $(FPIC) $(CCFLAGS) $^ -o $@
 
-# Fake up a linker script that will export *just* the PyInit entry
-# point we want. (If we don't do this we can have interesting failures
-# when loading multiple of these Python extensions in the same space.)
-ifeq ($(UNAME), Darwin)
-$(BIN)/ext/%.ldscript:
+# We take the native-code output of the Generator, add the Python-Extension
+# code (to make it callable from Python), and put the resulting Python AOT Extension
+# into the 'aot' folder.
+$(BIN)/aot/%.so: $(BIN)/%.py.o $(BIN)/%.a $(BIN)/runtime.a $(BIN)/%.ldscript
 	@echo Building $@...
 	@mkdir -p $(@D)
-	@echo _PyInit_$* > $@
+	@$(CXX) $(LDFLAGS) $(filter-out %.ldscript,$^) -shared $(subst %LDSCRIPT%,$(BIN)/$*.ldscript,$(PYEXT_LDSCRIPT_FLAG)) -o $@
 
-PYEXT_LDSCRIPT_FLAG = -Wl,-exported_symbols_list %LDSCRIPT%
-else
-# Assume Desktop Linux
-$(BIN)/ext/%.ldscript:
+# OK, now we want to produce a Stub Extension for the same Generator:
+# Compiling PyStub.cpp, then linking with the generator's .o file, PyStubImpl.o, plus the same libHalide
+# being used by halide.so.
+#
+# Note that we set HALIDE_PYSTUB_MODULE_NAME to $*_stub (e.g. foo_stub) but
+# set HALIDE_PYSTUB_GENERATOR_NAME to the unadorned name of the Generator.
+$(BIN)/%_PyStub.o: $(ROOT_DIR)/stub/PyStub.cpp
 	@echo Building $@...
 	@mkdir -p $(@D)
-	@echo "{" > $@
-	@echo "  global: PyInit_$*;" >> $@
-	@echo "  local: *;" >> $@
-	@echo "};" >> $@
-PYEXT_LDSCRIPT_FLAG = -Wl,--version-script=%LDSCRIPT%
-endif
+	@$(CXX) $(CCFLAGS) -DHALIDE_PYSTUB_MODULE_NAME=$*_stub -DHALIDE_PYSTUB_GENERATOR_NAME=$* -c $< -o $@
 
-# The Python extension of the generator is already in $(BIN), and is named
-# the same, so put the Python extension of the function into ext/.
-$(BIN)/ext/%.so: $(BIN)/%.py.o $(BIN)/%.a $(BIN)/runtime.a $(BIN)/ext/%.ldscript
+$(BIN)/stub/%_stub.so: $(BIN)/%_PyStub.o $(BIN)/PyStubImpl.o $(BIN)/%_generator.o $(BIN)/%_stub.ldscript $(LIBHALIDE)
 	@echo Building $@...
 	@mkdir -p $(@D)
-	@$(CXX) $(LDFLAGS) $(filter-out $(BIN)/ext/$*.ldscript,$^) -shared $(subst %LDSCRIPT%,$(BIN)/ext/$*.ldscript,$(PYEXT_LDSCRIPT_FLAG)) -o $@
+	@$(CXX) $(LDFLAGS) $(filter-out %.ldscript,$^) -shared $(subst %LDSCRIPT%,$(BIN)/$*_stub.ldscript,$(PYEXT_LDSCRIPT_FLAG)) -o $@
 
-test_correctness_addconstant_test: $(BIN)/ext/addconstant.so
-test_correctness_bit_test: $(BIN)/ext/bit.so
-test_correctness_user_context_test: $(BIN)/ext/user_context.so
-test_correctness_pystub: $(BIN)/generators/simplestub.so $(BIN)/generators/complexstub.so
+GENERATOR_SRCS=$(shell ls $(ROOT_DIR)/correctness/generators/*_generator.cpp)
+GENERATOR_AOT_EXTENSIONS=$(GENERATOR_SRCS:$(ROOT_DIR)/correctness/generators/%_generator.cpp=$(BIN)/aot/%.so)
+GENERATOR_STUB_EXTENSIONS=$(GENERATOR_SRCS:$(ROOT_DIR)/correctness/generators/%_generator.cpp=$(BIN)/stub/%_stub.so)
 
 APPS = $(shell ls $(ROOT_DIR)/apps/*.py)
 CORRECTNESS = $(shell ls $(ROOT_DIR)/correctness/*.py)
@@ -182,10 +217,13 @@ test_apps_%: $(ROOT_DIR)/apps/%.py $(MODULE)
 .PHONY: test_correctness
 test_correctness: $(CORRECTNESS:$(ROOT_DIR)/correctness/%.py=test_correctness_%)
 
-test_correctness_%: $(ROOT_DIR)/correctness/%.py $(MODULE)
+# For simplicity of the build system, we just have every correctness test depend
+# on every Generator and AOT extension. Normally this would not be a good idea,
+# but it's fine for us here.
+test_correctness_%: $(ROOT_DIR)/correctness/%.py $(MODULE) $(GENERATOR_AOT_EXTENSIONS) $(GENERATOR_STUB_EXTENSIONS)
 	@echo Testing $*...
 	@mkdir -p $(TEST_TMP)
-	@cd $(TEST_TMP); PYTHONPATH="$(BIN)/ext:$(BIN)/generators:$(BIN):$$PYTHONPATH" $(PYTHON) $<
+	@cd $(TEST_TMP); PYTHONPATH="$(BIN)/aot:$(BIN)/stub:$(BIN):$$PYTHONPATH" $(PYTHON) $<
 
 .PHONY: test_tutorial
 test_tutorial: $(TUTORIAL:$(ROOT_DIR)/tutorial/%.py=test_tutorial_%)

python_bindings/correctness/CMakeLists.txt

@@ -1,15 +1,4 @@
-set(GENERATORS
-    complexstub_generator.cpp
-    simplestub_generator.cpp
-    )
-
-foreach (GEN IN LISTS GENERATORS)
-    string(REPLACE "_generator.cpp" "" TARGET "${GEN}")
-    add_generator_python(${TARGET} ${GEN})
-endforeach ()
-
-# Handle addconstant, bit, user_context
-add_subdirectory(ext)
+add_subdirectory(generators)
 
 add_library(the_sort_function MODULE the_sort_function.c)
 target_link_libraries(the_sort_function PRIVATE Halide::Runtime)
@@ -38,12 +27,12 @@ set(TESTS
     var.py
     )
 
-# Use generator expressions to get the true output paths of these files
-# CMAKE_CURRENT_BINARY_DIR is incorrect.
+# Use generator expressions to get the true output paths of these files.
 make_shell_path(PYTHONPATH
-                "$<TARGET_FILE_DIR:py_addconstant>"
-                "$<TARGET_FILE_DIR:simplestub>"
-                "$<TARGET_FILE_DIR:Halide::Python>")
+                "$<TARGET_FILE_DIR:py_aot_bit>"
+                "$<TARGET_FILE_DIR:py_stub_bit>"
+                "$<TARGET_FILE_DIR:Halide::Python>"
+                )
 
 foreach (TEST IN LISTS TESTS)
     get_filename_component(TEST_NAME ${TEST} NAME_WE)