Possibility to align a split.

[mcourteaux]

Opening PR to show work in progress, and gather feedback and enter discussion.

Instead of splitting a loop where the rewritten inner and outer loop always start at 0, like this:

// f.split(x, xo, xi, 32);
for (xo, 0, (f.extent.0 + 31) / 32 - 1) {
  for (xi, 0, 31) {
      let x = xo * 32 + xi 
  }
}

It's now possible to align the first iteration of the inner loop, like so:

f.split(x, xo, xi, 32, -4);
for (xo, -4 / 32, (f.extent.0 - 4 + 31) / 32 - 1) {
  for (xi, 0 - 4, 31 - 4) {
      let x = xo * 32 + xi
      if (x >= 0 && x <= max) { // for GuardWithIf

      }
  }
}

This allows you to then unroll or vectorize the inner loop with a known alignment (modulo).
This comes up when demosaicing Bayer images where the offset of the filter pattern (CFA) is not known up front. Instead of compiling 4 different specializations of this pipeline with all for possible offsets, you can now pass in the CFA-offset as a runtime Param<int>:

f(x, y, c) = select((x + offset_x) % 2 == 0, /* similar select for (y + y_offset) */); 
f.split(x, xo, xi, 2, offset_x)
 .split(y, yo, yi, 2, offset_y)
 .reorder(c, xi, yi, xo, yo)
 .unroll(xi)
 .unroll(yi);

Alternative considered

After a very lengthy discussion with @abadams I attempted to implement a .guard_with_if() directive that would combine orthogonally with .align_bounds(). However, align bounds changes the bounds during bounds inference phase. The initial idea of "fixing" the widend bounds was to protect it with an if (hence guard_with_if()). However, what the if is supposed to do is to guard against out-of-bounds accesses, but align_bounds() actually changes the bounds, so there is nothing to protect against. For more details, see #9357.

While implementing this, the number of things that broke, missing simplifier rules to make it work, new behavior required in BoundsInference, BoundConstantExtentLoops, SlidingWindow was not pretty. While I think I got all of this working correctly in mcourteaux/guard-with-if, it is fundamentally backwards-incompatible because .align_bounds() is meant to change the size of the bounds (and impose constraints on it): it extends the computed region. Combining it with a ShiftInwards, or guard_with_if() again shrinks the computed region. There were several tests depending on the widening behavior of align_bounds().

Breaking changes

None, it's a new feature that isn't used anywhere.
If we can agree this is a good idea, I'll keep working on this to get it with the necessary documentation, tutorial, serialization, python bindings.

Checklist

As I said, checklist to be completed if others greenlight this.

• Tests added or updated (not required for docs, CI config, or typo fixes)
• Documentation updated (if public API changed)
• Python bindings updated (if public API changed)
• Benchmarks are included here if the change is intended to affect performance.
• Commits include AI attribution where applicable (see Code of Conduct)

[abadams]

Could you provide an example where the starting loop is not already aligned in absolute coordinates? In your example before case it's aligned to zero.

Also, how do aligned splits interact with rfactor? It has to enact any relevant splits eagerly. Hopefully rfactor tolerates this as-is, but some adding some rfactor test cases that use aligned splits is probably a good idea.

[mcourteaux]

Could you provide an example where the starting loop is not already aligned in absolute coordinates?

The starting loop is aligned in absolute coordinates. The problem is that the split() always rebases the inner and outer loop to zero. So your original loop was aligned (before rebase loops to zero lowering pass), but the initial scheduling rewrites the aligned loop by a split into two non-aligned loops. That happens here (notice + old_min):

Halide/src/ApplySplit.cpp

Line 29 in 5c21c82

Expr base = outer * split.factor + old_min;

and

Halide/src/ApplySplit.cpp

Lines 176 to 181 in 5c21c82

	Expr inner_extent = split.factor;
	Expr outer_extent = (old_var_max - old_var_min + split.factor) / split.factor;
	let_stmts.emplace_back(prefix + split.inner + ".loop_min", 0);
	let_stmts.emplace_back(prefix + split.inner + ".loop_max", inner_extent - 1);
	let_stmts.emplace_back(prefix + split.outer + ".loop_min", 0);
	let_stmts.emplace_back(prefix + split.outer + ".loop_max", outer_extent - 1);

[mcourteaux]

Still working on this. It's taking more and more shape. Will update tomorrow most likely. I ran into a really nice use case for combining align_bounds() on a producer, and an aligned split() on a consumer.

[abadams]

A use case I found today is roughly:

f(x) = select(x <= 0, ..., x < 7, likely(... x % 2 ...), ...)
f.bound(x, 0, 8)
f.split(x, xo, xi, 2, 1, TailStrategy::GuardWithIf).always_partition(xo);

I.e. I have a loop of size 8, and I want to unroll it 2x to remove a modulo, but I also want to carve off the first and last iteration as different. Does this make sense as a use of the feature? Partitioning xo without the alignment today carves off the first two and last two iterations, which is not what I want.

[mcourteaux]

Does this make sense as a use of the feature? Partitioning xo without the alignment today carves off the first two and last two iterations, which is not what I want.

Cool! Yep, I think it fully makes sense. I'll add this particular case as a test to make sure it does what we expect here.

The test works!

 produce f {
  f[0] = 100
  for (f.s0.x.xo, 0, 2) {
   f[(f.s0.x.xo*2) + 1] = 1
   f[(f.s0.x.xo*2) + 2] = 0
  }
  f[7] = 200
 }

with this schedule indeed:

    Var x{"x"}, xo{"xo"}, xi{"xi"};
    Func f{"f"};
    f(x) = select(x <= 0, 100,
                  x < 7, likely(x % 2),
                  200);
    f.bound(x, 0, 8);
    f.split(x, xo, xi, 2, 1, TailStrategy::GuardWithIf)
        .always_partition(xo)
        .unroll(xi);