Calling realize() within an AOT Generator is problematic

[steven-johnson]

A user inside Google recently had some confusion about unexpected results from a Generator they wrote; the issue turned out to be that some Halide-helper-library code that they were indirectly referencing called realize() to produce a LUT, and while this call referenced an Input<int> to the Generator, the LUT was (apparently) quietly produced with the value of zero.

This was an easy workaround (ie, use compute_root() instead of realize()), but there are at least two issues here worth investigating:

• Should it be an error to call .realize() (or any similar JIT-invocation) during AOT Generator evaluation? While I can envision not-reasonable use cases (such as LUT creation, as in this case), they all seem like corner cases that have better solutions in the first plce (such as using compute_root()).
• Should it be possible to JIT-compile with a Parameter that doesn't have an explicit value set? (For ImageParam I suspect we don't allow this, but perhaps Param<> is just quietly defaulting to a value of zero?) I'd argue that the answer should be "no" in all cases.

[zvookin]

Calling realize within an AOT generator is a standard technique. I used to cover it in presentations I did on how to use Halide. Moreover, in as much as Halide provides JIT as an API, not being able to JIT in this context would signal something broken elsewhere in Halide.

Does the threadsafe Callable mechanism make things better by requiring all the arguments get passed in at the call site?

[abadams]

I also think we want to allow realize within AOT generators.

It might be that we never track if a scalar param has been set, and are just defaulting to zero. We should probably add a flag for that so that we can throw an error, like we do for buffers.

[steven-johnson]

Fixing this for Param<> in the JIT case is indeed pretty trivial (see #7853); as you suggested, adding a flag to track it works fine, because when

However, this does not fix it for the Generator/AOT case (which is how we discovered it in the first place); the failure mode we saw was basically something like this:

class ShouldFail : public Halide::Generator<ShouldFail> {
public:
    Input<Buffer<uint8_t, 1>> input{"input"};
    Input<float> scalar_input{"scalar_input"};
    Output<Buffer<uint8_t, 1>> output{"output"};

    void generate() {
        Func lut_fn("lut_fn");
        lut_fn(x) = u8_sat(x * scalar_input / 255.f);

        // This should always fail, because it depends on a scalar input
        // that *cannot* have a valid value at this point.
        auto lut = lut_fn.realize({256});

        output(x) = input(x) + lut[0](x);
    }
};

Obviously, scalar_input has no defined value, but this doesn't fail as-is or with #7853 in place. The problem here is that there isn't actually any Parameter involved in the realize call: the Expr that is returned by accessing the scalar_input member field here is a Variable (specifically, Variable::make("scalar_input") in this case). When we call realize(), there aren't any Parameter values involved, so the newly added check is never made. When we get into CodeGenLLVM::compile() and create an llvm::Function, we (correctly) deduce that this should be an input to the func... but apparently the default value that gets put in the symbol table for the name is whatever LLVM has as the default value for an arg of that type, which is presumably zero, which is what we use.

At this point it's tricky to figure out we are referencing a bogus Parameter, because we have gotten here by getting a list of LoweredFunc and looking at LoweredArgument, which came from Argument, which came from the Parameter, but (not unreasonably) the info about jit-specific uninited values is lost in the process.

Not sure the right approach to fix this. A few thoughts:

• Add a flag to Argument that is propagated from Parameter? Probably would work but feels janky and weird.
• Add a visitor to the JIT code that looks for Variable nodes that have no Let/LetStmt? This might be the better answer, as it is localized entirely to within the JIT code instead of leaking out elsewhere, but I feel like the correct fix is probably more complicated.
• Something else?

[abadams]

We shouldn't be making a raw Variable when we use scalar_input as an Expr. Variables are supposed to either be bound by an enclosing Let or For, or have a Parameter associated with them. We should probably validate this in the IR at some point. I'm a bit confused about how more things aren't breaking if we have free Variables just floating around. I guess most lowering passes are coded defensively.

[abadams]

Wait, I think I'm not understanding. The only Variable::make I see in Generator.cpp attaches a Parameter object.

[steven-johnson]

Wait, I think I'm not understanding. The only Variable::make I see in Generator.cpp attaches a Parameter object.

Let me take another look, maybe I'm confused

[steven-johnson]

We shouldn't be making a raw Variable when we use scalar_input as an Expr. Variables are supposed to either be bound by an enclosing Let or For, or have a Parameter associated with them

FYI, it was in fact a Variable-with-Parameter-associated (not a 'raw' Variable), I misread the code