[serialization] Add serialization support to generator interface (hal…

…ide#7792) * Add serialization support to Generator interface * Clang format pass * Make target required when emitting a serialized pipeline (since schedule may be target dependent). Apply auto-scheduler before serialization so that schedules can be serialized. * Fix enum ordering for hlpipe. Fix hlpipe comments. Add missing hlpipe enum to pyenums. * Remove unused Serialization build_mode * Fix formatting * Remove unused serializable flag. Remove redundant cpp_stub check. Fix comments. * Safeguard emit_hlpipe calls with #ifdef WITH_SERIALIZATION --------- Co-authored-by: Derek Gerstmann <dgerstmann@adobe.com> Co-authored-by: Steven Johnson <srj@google.com>
ardier · Mar 3, 2024 · 7d2e9d9 · 7d2e9d9
1 parent 33e925a
commit 7d2e9d9
Show file tree

Hide file tree

Showing 9 changed files with 98 additions and 20 deletions.
diff --git a/python_bindings/src/halide/halide_/PyEnums.cpp b/python_bindings/src/halide/halide_/PyEnums.cpp
@@ -205,6 +205,7 @@ void define_enums(py::module &m) {
         .value("cpp_stub", OutputFileType::cpp_stub)
         .value("featurization", OutputFileType::featurization)
         .value("function_info_header", OutputFileType::function_info_header)
+        .value("hlpipe", OutputFileType::hlpipe)
         .value("llvm_assembly", OutputFileType::llvm_assembly)
         .value("object", OutputFileType::object)
         .value("python_extension", OutputFileType::python_extension)

diff --git a/python_bindings/src/halide/halide_/PyGenerator.cpp b/python_bindings/src/halide/halide_/PyGenerator.cpp
@@ -100,6 +100,12 @@ class PyGeneratorBase : public AbstractGenerator {
         // but don't throw an error, just return false.
         return false;
     }
+
+    bool emit_hlpipe(const std::string & /*hlpipe_file_path*/) override {
+        // Python Generators don't support this yet ...
+        // but don't throw an error, just return false.
+        return false;
+    }
 };
 
 class PyGeneratorFactoryProvider : public GeneratorFactoryProvider {

diff --git a/src/AbstractGenerator.h b/src/AbstractGenerator.h
@@ -177,6 +177,21 @@ class AbstractGenerator {
      */
     virtual bool emit_cpp_stub(const std::string &stub_file_path) = 0;
 
+    /** Emit a Serialized Halide Pipeline (.hlpipe) file to the given path. Not all Generators support this.
+     *
+     * If you call this method, you should not call any other AbstractGenerator methods
+     * on this instance, before or after this call.
+     *
+     * If the Generator is capable of emitting an hlpipe, do so and return true. (Errors
+     * during hlpipe emission should assert-fail rather than returning false.)
+     *
+     * If the Generator is not capable of emitting an hlpipe, do nothing and return false.
+     *
+     * CALL-AFTER: none
+     * CALL-BEFORE: none
+     */
+    virtual bool emit_hlpipe(const std::string &hlpipe_file_path) = 0;
+
     /** By default, a Generator must declare all Inputs before all Outputs.
      *  In some unusual situations (e.g. metaprogramming situations), it's
      * desirable to allow them to be declared out-of-order and put the onus

diff --git a/src/Generator.cpp b/src/Generator.cpp
@@ -10,6 +10,7 @@
 #include "Generator.h"
 #include "IRPrinter.h"
 #include "Module.h"
+#include "Serialization.h"
 #include "Simplify.h"
 
 #ifdef HALIDE_ALLOW_GENERATOR_BUILD_METHOD
@@ -651,7 +652,7 @@ gengen
  -e  A comma separated list of files to emit. Accepted values are:
      [assembly, bitcode, c_header, c_source, cpp_stub, featurization,
       llvm_assembly, object, python_extension, pytorch_wrapper, registration,
-      schedule, static_library, stmt, stmt_html, compiler_log].
+      schedule, static_library, stmt, stmt_html, compiler_log, hlpipe].
      If omitted, default value is [c_header, static_library, registration].
 
  -p  A comma-separated list of shared libraries that will be loaded before the
@@ -795,6 +796,7 @@ gengen
             std::map<std::string, OutputFileType> output_name_to_enum = {
                 {"cpp", OutputFileType::c_source},
                 {"h", OutputFileType::c_header},
+                {"hlpipe", OutputFileType::hlpipe},
                 {"html", OutputFileType::stmt_html},
                 {"o", OutputFileType::object},
                 {"py.c", OutputFileType::python_extension},
@@ -998,6 +1000,7 @@ void execute_generator(const ExecuteGeneratorArgs &args_in) {
 
     const bool cpp_stub_only = args.output_types.size() == 1 &&
                                args.output_types.count(OutputFileType::cpp_stub) == 1;
+
     if (!cpp_stub_only) {
         // It's ok to leave targets unspecified if we are generating *only* a cpp_stub
         internal_assert(!args.targets.empty());
@@ -1047,27 +1050,49 @@ void execute_generator(const ExecuteGeneratorArgs &args_in) {
     if (!args.generator_name.empty()) {
         const std::string base_path = args.output_dir + "/" + args.file_base_name;
         debug(1) << "Generator " << args.generator_name << " has base_path " << base_path << "\n";
+
+        // Factory function for creating a generator instance for a given function name and target
+        auto generator_factory = [&](const std::string &function_name, const Target &target) -> AbstractGeneratorPtr {
+            // Must re-create each time since each instance will have a different Target.
+            auto gen = args.create_generator(args.generator_name, GeneratorContext(target));
+            for (const auto &kv : args.generator_params) {
+                if (kv.first == "target") {
+                    continue;
+                }
+                gen->set_generatorparam_value(kv.first, kv.second);
+            }
+            return gen;
+        };
+
         if (args.output_types.count(OutputFileType::cpp_stub)) {
-            // When generating cpp_stub, we ignore all generator args passed in, and supply a fake Target.
+            // When generating cpp_stub we ignore all generator args passed in, and supply a fake Target.
             // (CompilerLogger is never enabled for cpp_stub, for now anyway.)
             const Target fake_target = Target();
             auto gen = args.create_generator(args.generator_name, GeneratorContext(fake_target));
             auto output_files = compute_output_files(fake_target, base_path, args.output_types);
             gen->emit_cpp_stub(output_files[OutputFileType::cpp_stub]);
         }
 
+#ifdef WITH_SERIALIZATION
+        if (args.output_types.count(OutputFileType::hlpipe)) {
+            // When serializing a halide pipeline, target is required (since the schedule may be target dependent).
+            // If multiple targets are specified, add the target name as a suffix to the filename.
+            const bool use_target_suffix = (args.targets.size() > 1);
+            for (size_t i = 0; i < args.targets.size(); ++i) {
+                const Target &target = args.targets[i];
+                const std::string hlpipe_path = use_target_suffix ? (base_path + "-" + target.to_string()) : base_path;
+                auto output_files = compute_output_files(target, hlpipe_path, args.output_types);
+                auto gen = generator_factory(args.function_name, target);
+                gen->emit_hlpipe(output_files[OutputFileType::hlpipe]);
+            }
+        }
+#endif
+
         // Don't bother with this if we're just emitting a cpp_stub.
         if (!cpp_stub_only) {
             auto output_files = compute_output_files(args.targets[0], base_path, args.output_types);
             auto module_factory = [&](const std::string &function_name, const Target &target) -> Module {
-                // Must re-create each time since each instance will have a different Target.
-                auto gen = args.create_generator(args.generator_name, GeneratorContext(target));
-                for (const auto &kv : args.generator_params) {
-                    if (kv.first == "target") {
-                        continue;
-                    }
-                    gen->set_generatorparam_value(kv.first, kv.second);
-                }
+                auto gen = generator_factory(function_name, target);
                 return args.build_mode == ExecuteGeneratorArgs::Gradient ?
                            gen->build_gradient_module(function_name) :
                            gen->build_module(function_name);
@@ -1609,6 +1634,26 @@ bool GeneratorBase::emit_cpp_stub(const std::string &stub_file_path) {
     return true;
 }
 
+bool GeneratorBase::emit_hlpipe(const std::string &hlpipe_file_path) {
+#ifdef WITH_SERIALIZATION
+    user_assert(!generator_registered_name.empty() && !generator_stub_name.empty()) << "Generator has no name.\n";
+    Pipeline pipeline = build_pipeline();
+    AutoSchedulerResults auto_schedule_results;
+    const auto context = this->context();
+    const auto &asp = context.autoscheduler_params();
+    if (!asp.name.empty()) {
+        debug(1) << "Applying autoscheduler " << asp.name << " to Generator " << name() << " ...\n";
+        auto_schedule_results = pipeline.apply_autoscheduler(context.target(), asp);
+    }
+    std::map<std::string, Internal::Parameter> params;  // FIXME: Remove when API allows this to be optional
+    serialize_pipeline(pipeline, hlpipe_file_path, params);
+    return true;
+#else
+    user_error << "Serialization is not supported in this build of Halide; try rebuilding with WITH_SERIALIZATION=ON.";
+    return false;
+#endif
+}
+
 GIOBase::GIOBase(size_t array_size,
                  const std::string &name,
                  ArgInfoKind kind,

diff --git a/src/Generator.h b/src/Generator.h
@@ -3689,6 +3689,7 @@ class GeneratorBase : public NamesInterface, public AbstractGenerator {
     void bind_input(const std::string &name, const std::vector<Expr> &v) override;
 
     bool emit_cpp_stub(const std::string &stub_file_path) override;
+    bool emit_hlpipe(const std::string &hlpipe_file_path) override;
 
     GeneratorBase(const GeneratorBase &) = delete;
     GeneratorBase &operator=(const GeneratorBase &) = delete;
@@ -3921,7 +3922,7 @@ struct ExecuteGeneratorArgs {
         Default,
 
         // Build a version suitable for using for gradient descent calculation.
-        Gradient
+        Gradient,
     } build_mode = Default;
 
     // The fn that will produce Generator(s) from the name specified.

diff --git a/src/Module.cpp b/src/Module.cpp
@@ -40,6 +40,7 @@ std::map<OutputFileType, const OutputInfo> get_output_info(const Target &target)
         {OutputFileType::cpp_stub, {"cpp_stub", ".stub.h", IsSingle}},
         {OutputFileType::featurization, {"featurization", ".featurization", IsMulti}},
         {OutputFileType::function_info_header, {"function_info_header", ".function_info.h", IsSingle}},
+        {OutputFileType::hlpipe, {"hlpipe", ".hlpipe", IsSingle}},
         {OutputFileType::llvm_assembly, {"llvm_assembly", ".ll", IsMulti}},
         {OutputFileType::object, {"object", is_windows_coff ? ".obj" : ".o", IsMulti}},
         {OutputFileType::python_extension, {"python_extension", ".py.cpp", IsSingle}},

diff --git a/src/Module.h b/src/Module.h
@@ -32,6 +32,7 @@ enum class OutputFileType {
     cpp_stub,
     featurization,
     function_info_header,
+    hlpipe,
     llvm_assembly,
     object,
     python_extension,

diff --git a/test/correctness/compile_to_multitarget.cpp b/test/correctness/compile_to_multitarget.cpp
@@ -192,15 +192,18 @@ void test_compile_to_everything(Func j, bool do_object) {
         // {OutputFileType::cpp_stub, filename_prefix + ".stub.h"},  // IsSingle
         {OutputFileType::featurization, filename_prefix + ".featurization"},           // IsMulti
         {OutputFileType::function_info_header, filename_prefix + ".function_info.h"},  // IsSingle
-        {OutputFileType::llvm_assembly, filename_prefix + ".ll"},                      // IsMulti
-        {OutputFileType::object, filename_prefix + o},                                 // IsMulti
-        {OutputFileType::python_extension, filename_prefix + ".py.cpp"},               // IsSingle
-        {OutputFileType::pytorch_wrapper, filename_prefix + ".pytorch.h"},             // IsSingle
-        {OutputFileType::registration, filename_prefix + ".registration.cpp"},         // IsSingle
-        {OutputFileType::schedule, filename_prefix + ".schedule.h"},                   // IsSingle
-        {OutputFileType::static_library, filename_prefix + a},                         // IsSingle
-        {OutputFileType::stmt, filename_prefix + ".stmt"},                             // IsMulti
-        {OutputFileType::stmt_html, filename_prefix + ".stmt.html"},                   // IsMulti
+        // Note: compile_multitarget() doesn't produce hlpipe output,
+        // even if you pass this in (since the pipeline has already been lowered)
+        // {OutputFileType::hlpipe, filename_prefix + ".hlpipe"},  // IsSingle
+        {OutputFileType::llvm_assembly, filename_prefix + ".ll"},               // IsMulti
+        {OutputFileType::object, filename_prefix + o},                          // IsMulti
+        {OutputFileType::python_extension, filename_prefix + ".py.cpp"},        // IsSingle
+        {OutputFileType::pytorch_wrapper, filename_prefix + ".pytorch.h"},      // IsSingle
+        {OutputFileType::registration, filename_prefix + ".registration.cpp"},  // IsSingle
+        {OutputFileType::schedule, filename_prefix + ".schedule.h"},            // IsSingle
+        {OutputFileType::static_library, filename_prefix + a},                  // IsSingle
+        {OutputFileType::stmt, filename_prefix + ".stmt"},                      // IsMulti
+        {OutputFileType::stmt_html, filename_prefix + ".stmt.html"},            // IsMulti
     };
     if (do_object) {
         outputs.erase(OutputFileType::static_library);

diff --git a/test/generator/abstractgeneratortest_generator.cpp b/test/generator/abstractgeneratortest_generator.cpp
@@ -135,6 +135,11 @@ class AbstractGeneratorTest : public AbstractGenerator {
         // not supported
         return false;
     }
+
+    bool emit_hlpipe(const std::string & /*hlpipe_file_path*/) override {
+        // not supported
+        return false;
+    }
 };
 
 RegisterGenerator register_something(AbstractGeneratorTestName,