Stage

class Stage

A single definition of a Func.

May be a pure or update definition.

Public Functions

inline const Internal::StageSchedule &get_schedule() const

Return the current StageSchedule associated with this Stage.

For introspection only: to modify schedule, use the Func interface.

std::string dump_argument_list() const

Return a string describing the current var list taking into account all the splits, reorders, and tiles.

std::string name() const

Return the name of this stage, e.g.

“f.update(2)”

Func rfactor(const std::vector<std::pair<RVar, Var>> &preserved)

Calling rfactor() on an associative update definition a Func will split the update into an intermediate which computes the partial results and replaces the current update definition with a new definition which merges the partial results.

If called on a init/pure definition, this will throw an error. rfactor() will automatically infer the associative reduction operator and identity of the operator. If it can’t prove the operation is associative or if it cannot find an identity for that operator, this will throw an error. In addition, commutativity of the operator is required if rfactor() is called on the inner dimension but excluding the outer dimensions.

rfactor() takes as input ‘preserved’, which is a list of <RVar, Var> pairs. The rvars not listed in ‘preserved’ are removed from the original Func and are lifted to the intermediate Func. The remaining rvars (the ones in ‘preserved’) are made pure in the intermediate Func. The intermediate Func’s update definition inherits all scheduling directives (e.g. split,fuse, etc.) applied to the original Func’s update definition. The loop order of the intermediate Func’s update definition is the same as the original, although the RVars in ‘preserved’ are replaced by the new pure Vars. The loop order of the intermediate Func’s init definition from innermost to outermost is the args’ order of the original Func’s init definition followed by the new pure Vars.

The intermediate Func also inherits storage order from the original Func with the new pure Vars added to the outermost.

For example, f.update(0).rfactor({{r.y, u}}) would rewrite a pipeline like this:

f(x, y) = 0;
f(x, y) += g(r.x, r.y);

into a pipeline like this:

f_intm(x, y, u) = 0;
f_intm(x, y, u) += g(r.x, u);

f(x, y) = 0;
f(x, y) += f_intm(x, y, r.y);

This has a variety of uses. You can use it to split computation of an associative reduction:

f(x, y) = 10;
RDom r(0, 96);
f(x, y) = max(f(x, y), g(x, y, r.x));
f.update(0).split(r.x, rxo, rxi, 8).reorder(y, x).parallel(x);
f.update(0).rfactor({{rxo, u}}).compute_root().parallel(u).update(0).parallel(u);

, which is equivalent to:

parallel for u = 0 to 11:
  for y:
    for x:
      f_intm(x, y, u) = -inf
parallel for x:
  for y:
    parallel for u = 0 to 11:
      for rxi = 0 to 7:
        f_intm(x, y, u) = max(f_intm(x, y, u), g(8*u + rxi))
for y:
  for x:
    f(x, y) = 10
parallel for x:
  for y:
    for rxo = 0 to 11:
      f(x, y) = max(f(x, y), f_intm(x, y, rxo))

Stage &eager_inline(const std::vector<Func> &fs)

Immediately inline direct calls to each of the given Funcs into this stage’s definition.

The Funcs are inlined in dependency order regardless of the order they are passed, so if one inlined Func’s body calls another, both are fully folded in.

Unlike compute_inline(), which merely marks a Func to be inlined during lowering, eager_inline() performs the substitution now, at schedule time, rewriting only this stage’s definition in place. This is useful to surface structure that other schedule-time directives (e.g. rfactor()) need to see.

Each inlined Func must be inlinable: a pure Func (no update or extern definition) with no specializations, and with a schedule compatible with inlining (as for compute_inline()). The inlined Funcs are otherwise unchanged; only this stage’s calls to them are replaced.

Stage &compute_with(LoopLevel loop_level, const std::vector<std::pair<VarOrRVar, LoopAlignStrategy>> &align)

Schedule the iteration over this stage to be fused with another stage ‘s’ from outermost loop to a given LoopLevel.

‘this’ stage will be computed AFTER ‘s’ in the innermost fused dimension. There should not be any dependencies between those two fused stages. If either of the stages being fused is a stage of an extern Func, this will throw an error.

Note that the two stages that are fused together should have the same exact schedule from the outermost to the innermost fused dimension, and the stage we are calling compute_with on should not have specializations, e.g. f2.compute_with(f1, x) is allowed only if f2 has no specializations.

Also, if a producer is desired to be computed at the fused loop level, the function passed to the compute_at() needs to be the “parent”. Consider the following code:

input(x, y) = x + y;
f(x, y) = input(x, y);
f(x, y) += 5;
g(x, y) = x - y;
g(x, y) += 10;
f.compute_with(g, y);
f.update().compute_with(g.update(), y);

To compute ‘input’ at the fused loop level at dimension y, we specify input.compute_at(g, y) instead of input.compute_at(f, y) since ‘g’ is the “parent” for this fused loop (i.e. ‘g’ is computed first before ‘f’ is computed). On the other hand, to compute ‘input’ at the innermost dimension of ‘f’, we specify input.compute_at(f, x) instead of input.compute_at(g, x) since the x dimension of ‘f’ is not fused (only the y dimension is).

Given the constraints, this has a variety of uses. Consider the following code:

f(x, y) = x + y;
g(x, y) = x - y;
h(x, y) = f(x, y) + g(x, y);
f.compute_root();
g.compute_root();
f.split(x, xo, xi, 8);
g.split(x, xo, xi, 8);
g.compute_with(f, xo);

This is equivalent to:

for y:
  for xo:
    for xi:
      f(8*xo + xi) = (8*xo + xi) + y
    for xi:
      g(8*xo + xi) = (8*xo + xi) - y
for y:
  for x:
    h(x, y) = f(x, y) + g(x, y)

The size of the dimensions of the stages computed_with do not have to match. Consider the following code where ‘g’ is half the size of ‘f’:

Image<int> f_im(size, size), g_im(size/2, size/2);
input(x, y) = x + y;
f(x, y) = input(x, y);
g(x, y) = input(2*x, 2*y);
g.compute_with(f, y);
input.compute_at(f, y);
Pipeline({f, g}).realize({f_im, g_im});

This is equivalent to:

for y = 0 to size-1:
  for x = 0 to size-1:
    input(x, y) = x + y;
  for x = 0 to size-1:
    f(x, y) = input(x, y)
  for x = 0 to size/2-1:
    if (y < size/2-1):
      g(x, y) = input(2*x, 2*y)

‘align’ specifies how the loop iteration of each dimension of the two stages being fused should be aligned in the fused loop nests (see LoopAlignStrategy for options). Consider the following loop nests:

for z = f_min_z to f_max_z:
  for y = f_min_y to f_max_y:
    for x = f_min_x to f_max_x:
      f(x, y, z) = x + y + z
for z = g_min_z to g_max_z:
  for y = g_min_y to g_max_y:
    for x = g_min_x to g_max_x:
      g(x, y, z) = x - y - z

If no alignment strategy is specified, the following loop nest will be generated:

for z = min(f_min_z, g_min_z) to max(f_max_z, g_max_z):
  for y = min(f_min_y, g_min_y) to max(f_max_y, g_max_y):
    for x = f_min_x to f_max_x:
      if (f_min_z <= z <= f_max_z):
        if (f_min_y <= y <= f_max_y):
          f(x, y, z) = x + y + z
    for x = g_min_x to g_max_x:
      if (g_min_z <= z <= g_max_z):
        if (g_min_y <= y <= g_max_y):
          g(x, y, z) = x - y - z

Instead, these alignment strategies:

g.compute_with(f, y, {{z, LoopAlignStrategy::AlignStart}, {y, LoopAlignStrategy::AlignEnd}});

will produce the following loop nest:

f_loop_min_z = f_min_z
f_loop_max_z = max(f_max_z, (f_min_z - g_min_z) + g_max_z)
for z = f_min_z to f_loop_max_z:
  f_loop_min_y = min(f_min_y, (f_max_y - g_max_y) + g_min_y)
  f_loop_max_y = f_max_y
  for y = f_loop_min_y to f_loop_max_y:
    for x = f_min_x to f_max_x:
      if (f_loop_min_z <= z <= f_loop_max_z):
        if (f_loop_min_y <= y <= f_loop_max_y):
          f(x, y, z) = x + y + z
    for x = g_min_x to g_max_x:
      g_shift_z = g_min_z - f_loop_min_z
      g_shift_y = g_max_y - f_loop_max_y
      if (g_min_z <= (z + g_shift_z) <= g_max_z):
        if (g_min_y <= (y + g_shift_y) <= g_max_y):
          g(x, y + g_shift_y, z + g_shift_z) = x - (y + g_shift_y) - (z + g_shift_z)

LoopAlignStrategy::AlignStart on dimension z will shift the loop iteration of ‘g’ at dimension z so that its starting value matches that of ‘f’. Likewise, LoopAlignStrategy::AlignEnd on dimension y will shift the loop iteration of ‘g’ at dimension y so that its end value matches that of ‘f’.

Stage &split(const VarOrRVar &old, const VarOrRVar &outer, const VarOrRVar &inner, const Expr &factor, TailStrategy tail = TailStrategy::Auto)

Scheduling calls that control how the domain of this stage is traversed.

See the documentation for Func for the meanings.

Stage &stream_stores()

Use non-temporal (streaming) stores for writes done by this Stage.

On targets that require it, Halide emits a fence immediately after this Stage’s production to ensure the streamed values are visible to whatever reads them next. Only legal on a Stage all of whose RVars (if any) are pure, i.e. already proven safe to parallelize: a Stage with a genuine loop-carried self-dependency could otherwise observe data it streamed earlier in the same Stage, before the fence. It is a user error to call this on a Stage that doesn’t meet this condition.

Stage &stream_loads()

Use non-temporal (streaming) loads for every direct read this Stage makes of another Func or external buffer (e.g.

an ImageParam). This is a hint to keep data that is only read once from displacing reusable data from the cache.

Stage &stream_loads(const std::vector<Func> &funcs)

Use non-temporal (streaming) loads for this Stage’s direct reads of the named Funcs.

It is a user error to name this Stage’s own Func (a self-load can’t be streamed).

std::vector<VarOrRVar> split_vars() const

Get the Vars and RVars of this definition, from innermost out, with splits applied.

This represents all the potentially-valid compute_at sites for this Stage. The RVars returned will be symbolic and not tied to a particular reduction domain, like the naked RVar objects used as split outputs. Note that this list by default will end with the sentinel Var::outermost.

void unscheduled()

Assert that this stage has intentionally been given no schedule, and suppress the warning about unscheduled update definitions that would otherwise fire.

This counts as a schedule, so calling this twice on the same Stage will fail the assertion.