Search code examples
halide

Can generators be used for JIT during development?


What is the best strategy for development in HALIDE? The end usage will be Ahead of time compiled using generators. Is there a way to invoke the functions defined in generator for JIT?

Thanks


Solution

  • Yes, Generators work just fine with JIT code.

    In general, Generators are the preferred way to encapsulate individual chunk of Halide for reuse. Until recently, they were a bit awkward to use in JIT mode, but recent changes that added machine-generated Stubs greatly simplify that and make using Generators in either JIT or AOT pretty simple.

    Unforunately, Generator Stubs are new enough that they aren't represented in the Tutorial yet; your best bet is to look at the self-tests in tests/generator for examples (specifically, example_jittest.cpp and stubtest_jittest.cpp).

    As an overview: The basic idea of a Generator Stub is that it's a machine-generated C++ class that is created based on the Generator's public description. It doesn't do anything that you couldn't do yourself directly, but it does make Generator usage much terser, simpler, and less error-prone.

    To produce a a Generator stub, just modify you your Makefile to add cpp_stub to the Generator's -e commandline flag , e.g.

    ./example.generator -n Example -o ./bin -e cpp_stub
    

    This will emit a C++ source file file named example.stub.h; inside this file you'll find a C++ class that looks something like this:

    class Example : public Halide::Internal::GeneratorStub {
    public:
      struct Inputs { 
        // One field per input to the Generator;
        // Buffer inputs will be Halide::Funcs, 
        // all other (scalar) inputs will be HalideExprs
      };
      struct GeneratorParams { 
        // One field per GeneratorParam in the Generator
      };
      struct ScheduleParams { 
        // One field per GeneratorParam in the Generator
      };
    
      Example();
    
      Example(
        const GeneratorContext* context,
        const Inputs& inputs,
        const GeneratorParams& params = GeneratorParams()
      );
    
      void schedule(const ScheduleParams& params = ScheduleParams());
    
      // Output(s)
      Func output;
      // If the Generator has multiple Outputs, they will be here too
    
    };
    

    You can use this Stub inside JIT code as though it were a helper function (well, mostly):

    #include "example.stub.h"
    
    Example::Inputs inputs = { ... };
    auto gen = Example(context, inputs);
    gen.schedule();
    Halide::Buffer<int32_t> img = gen.realize(kSize, kSize, 3);
    

    Under the hood, the Stub is constructing an instance of your Generator, filling in the Input parameters (based on the Inputs struct you give it), calling the generate() method to produce the Output Funcs, and returning them to you.

    Stub usage is woefully underdocumented as I type this; I've just opened an issue to gather up the documentation and examples we have into something that is more helpful to general Halide users.