Add special build for testing serialization via a serialization roundtrip in JIT compilation and fix serialization leaks

src/CMakeLists.txt

@@ -508,6 +508,15 @@ if (WITH_SERIALIZATION)
     target_compile_definitions(Halide PRIVATE WITH_SERIALIZATION)
 endif ()
 
+# Enable serialization testing by intercepting JIT compilation with a serialization roundtrip;
+# This is used only for special builds made specifically for testing, and must be disabled by default.
+option(WITH_SERIALIZATION_JIT_ROUNDTRIP_TESTING "Intercepting JIT compilation with a serialization roundtrip, for test only" OFF)
+if (WITH_SERIALIZATION_JIT_ROUNDTRIP_TESTING)
+    if (WITH_SERIALIZATION)
+        target_compile_definitions(Halide PRIVATE WITH_SERIALIZATION_JIT_ROUNDTRIP_TESTING)
+    endif ()
+endif ()
+
 add_library(Halide::Halide ALIAS Halide)
 
 target_link_libraries(Halide PRIVATE Halide::LLVM)

src/Pipeline.cpp

@@ -582,6 +582,25 @@ void Pipeline::compile_jit(const Target &target_arg) {
     // Clear all cached info in case there is an error.
     contents->invalidate_cache();
 
+#ifdef WITH_SERIALIZATION_JIT_ROUNDTRIP_TESTING
+    // TODO(https://github.com/halide/Halide/pull/7760): replace file serialization with in-memory serialization
+    std::string filename = generate_function_name() + ".hlpipe";
+    std::map<std::string, Internal::Parameter> external_params;
+    serialize_pipeline(*this, filename, external_params);
+    Pipeline deserialized_pipe = deserialize_pipeline(filename, external_params);
+    std::vector<Function> outputs;
+    for (const Func &f : deserialized_pipe.outputs()) {
+        outputs.push_back(f.function());
+    }
+    // We save the original output functions and requirements,
+    // and restore them once all lowering is done,
+    // so that reschedule/reorder storage can be properly handled.
+    std::vector<Function> origin_outputs = contents->outputs;
+    std::vector<Internal::Stmt> origin_requirements = contents->requirements;
+    contents->outputs = outputs;
+    contents->requirements = deserialized_pipe.requirements();
+#endif
+
     // Infer an arguments vector
     infer_arguments();
 
@@ -597,6 +616,11 @@ void Pipeline::compile_jit(const Target &target_arg) {
     Module module = compile_to_module(args, generate_function_name(), target).resolve_submodules();
     std::map<std::string, JITExtern> lowered_externs = contents->jit_externs;
     contents->jit_cache = compile_jit_cache(module, std::move(args), contents->outputs, contents->jit_externs, target);
+#ifdef WITH_SERIALIZATION_JIT_ROUNDTRIP_TESTING
+    // Restore the original outputs and requirements.
+    contents->outputs = origin_outputs;
+    contents->requirements = origin_requirements;
+#endif
 }
 
 Callable Pipeline::compile_to_callable(const std::vector<Argument> &args_in, const Target &target_arg) {