Reversing the compute_at/store_at relationship

[jrk]

Hypothesis: compute_at and store_at relationships should be defined at the consumer with respect to its producers, not at a producer with respect to its consumers. So (ignoring changing names to make the change clear), the schedule would say consumer.compute_at(producer, loop_var_of_consumer) rather than producer.compute_at(consumer, loop_var_of_consumer).

Context: @abadams and I were talking briefly about the remaining pain points / incidental complexity in using something like C++ classes in the vein of generators as the way to manage and structure Halide modules.

• Generators (or similar) work quite well for wrapping up and generating code from top-level pipelines.
• Generators (or similar) become more of a mess for modules which will be composed. IIRC there are two major issues:
  1. It gets verbose to write a long sequence of field assignments to wire up arguments. This is well addressed by stubs, which have hope of being reasonably approximated inline with some template work.
  2. Scheduling has to be done after algorithm generation, because in general it needs to have access to generator params of the downstream algorithm

Ignoring for a moment the obvious challenge that changing how we track and how the language specifies compute_at/store_at relationships, would this just solve (2)? It seems like it both makes module composition generally easier and reduces the number of stages (in the staged programming sense) from ~3 (generate algorithm, schedule algorithm, run pipeline) to ~2 (generate and schedule algorithm, run pipeline). That seems like a big mental and practical win.

It also seems like this model of producer-consumer relationships in scheduling (thinking from the consumer’s perspective how the producer is scheduled with respect to it) better matches the lowering process, anyway.

Thinking out loud, the main case it seems like this wouldn’t fully simplify is when you want to compute inner parts of a producer module relative to some outer context further downstream than the consumer module itself. Does this really come up? Still seems to require that fully general case allows deferred scheduling?

Also: does this mean compute_at should become a function on a module (i.e., on the standard Generator interface)? This would define the logical root for that module, which could be used just by the output Func, or potentially by many Funcs in the sub-pipeline. And this could be done now, as a standard convention for defining and setting up on or a few local-context loop levels as ScheduleParams, without changing Func, the innards of the schedule representation and lowering, or tens of thousands of lines of existing schedule code.

Another retrospectively obvious argument I realize in favor of this, and perhaps the most important both for having a clear mental model, and for improving modularity, is that it’s the consumer who knows the characteristics of its access, which determine the potential reuse and redundant compute of the producer.

[jrk]

/cc @zvookin @steven-johnson @dsharletg @shoaibkamil @psuriana

(I obviously realize this is a huge change, but I think it's potentially a big enough win that it's at least worth serious discussion.)

[steven-johnson]

This is well addressed by stubs, which have hope of being reasonably approximated inline with some template work

I have a prototype of what I think is a good approach (or at least an improvement over status quo) in https://github.com/halide/Halide/tree/srj-stub2 -- I haven't opened a PR on it yet because I think it still needs some thought. Opinions and suggestions welcomed.

[steven-johnson]

Scheduling has to be done after algorithm generation, because in general it needs to have access to generator params of the downstream algorithm

Not sure I understand this assertion -- can you give an example? The cases that I anticipated are more about gettting compute_at/store_at correct for intermediate calculations (which seems to be the intent of this Issue).

would this just solve (2)?

Either I'm missing something, or you left out some text before this sentence -- it's not at all clear to me what the 'this' is referring to.

[jrk]

Sorry Steven, you're right. Updated to clarify. Several of us have bounced this idea around a few times over the years, but I didn't initially realize that the only statement of the idea was in the title.

[steven-johnson]

the schedule would say consumer.compute_at(producer, loop_var_of_consumer) rather than producer.compute_at(consumer, loop_var_of_consumer)

How does producer.compute_root() change in this scenario?

Does LoopLevel usage change in this scenario? i.e., currently, it's equivalent to say

producer.compute_at(consumer, loop_var_of_consumer)

and

LoopLevel at(consumer, loop_var_of_consumer);
producer.compute_at(at);

[steven-johnson]

compute_at should become a function on a module
...
could be done now, as a standard convention for defining and setting up on or a few local-context loop levels as ScheduleParams

Yes, could be prototyped as a ScheduleParam already. Thinking ahead, though: we could simplify this further by adding an implicit ScheduleParam<LoopLevel> to each Output<>; these would default to LoopLevel::root() for toplevel compilation.

[jrk]

How does producer.compute_root() change in this scenario?

I think it would be consumer.compute_at(producer, LoopLevel::root()).

Does LoopLevel usage change in this scenario? i.e., currently, it's equivalent to say…

I think the second bit of example code here got broken on copy/paste?

[steven-johnson]

Does LoopLevel usage change in this scenario? i.e., currently, it's equivalent to say

producer.compute_at(consumer, loop_var_of_consumer)

and

LoopLevel at(consumer, loop_var_of_consumer);
producer.compute_at(at);

[jrk]

@abadams pointed out that a relatively similar convention was used inside [a commercial camera pipeline] and the FFT, but that it broke down at some point. @dsharletg do you remember why?

It's definitely true that some things can only be expressed with a global view, and not just per-module (with a producer module scheduling its own innards, and a consumer module scheduling its inputs). In particular, given:

f -> g -> h

you may still want to compute f at some level of h which is outside g outermost but below root. This definitely does require having a view from h through g to also include f at scheduling time, but I would still argue that the major concerns which lead to making this choice are still specific to their collective use inside the context of h, so that is a logical point at which to be specifying the schedule choices.

[jrk]

Also, to be clear: I don't believe that this changes the expressive power or semantics of the language, but that making clear and explicit a standard convention that consumers schedule the granularity at which to compute their producers, not the other way around, is likely the right model. This particularly relates to modular composition (e.g., as done in Halogen and Generators).

[zvookin]

I'll try to comment in more detail, but I don't see how this would actually solve any of the difficulties of writing components as all the same values have to be available in the same place. It is merely the syntax of the scheduling that changes. We've already looked somewhat at which things get scheduled inside a component and which are to be scheduled by its consumer (or "caller" more informally). E.g. generally intermediates are scheduled by the components and outputs of the component are scheduled by the consumer. This heuristic runs into issues because often the scheduling of intermediates is dependent on the scheduling of the output. E.g. in order to schedule an intermediate on an output, the output cannot be inline.

I think to really make progress on the use of the language, we are going to have to consider scheduling directives that can "skip over" Funcs in some circumstances, e.g. skip over ones that are inlined into exactly one place or something like that. But it is not obvious to me how to design these changes.

[jrk]

@zvookin I was just thinking about exactly this, below. We can talk offline today, too. I mostly wanted to write down some thoughts and seed more serious conversation about this.

Another detail that would arise here: it would have to be legal for a producer to specify its "domain order" (its own loops) even if its consumer decided it should be inlined. Maybe this is the sticking point Dillon encountered?

That works for a single producer Func in isolation, but the loops over the domain of the output Func of a producer module are also used for local compute_at/store_at granularity choices for earlier Funcs within the producer module. If I, say, locally tile and compute at tiles within the producer module, but then the consumer computes the producer output inline, what does this mean? Should this be an error?

If inline was identical to innermost, then this would work but the corresponding tile loops would simplify away to size 1 (assuming it was being computed at that granularity because the consumer just accesses it pointwise). Inline could perhaps be implicitly promoted to innermost in these cases, but that isn’t necessarily what you want.

Instead, it seems like what you often really want when you compute one output Func inline in a consumer is to transfer what would have been the module output Func’s schedule downstream through to the composed consumer Func. In effect, composition with pointwise dependence where you want to inline the output into its immediate use is really adding extra logic to the end of the upstream module, but would usually otherwise leave the producer module’s own choices (e.g., tiling wrt. its output) the same.

[jrk]

Talking to @dsharletg, [commercial camera pipeline] largely uses this pattern: producer modules locally schedule themselves as possible, and do their own domain order choices, while their consumers make the compute_at/store_at granularity choices for them. There seem to be two major gotchas:

1. Particularly for things scheduled inline into a consumer, you want their local scheduling choices to both still be valid/not cause an error, and ideally also propagate through to the consumer into which they’re inlined. This is discussed immediately above.
2. Using Var::outermost() works well for locally scheduling the intermediate computation within a module in a manner that’s flexible to be tiled with respect to downstream consumers, but it is not sufficient to line-buffer the intermediate stages with respect to the downstream consumer.