Overhaul of splitting, tail strategies, and loop partitioning.

[mcourteaux]

API

Stage::guard_with_if(VarOrRVar var);
Stage::guard_with_if(const std::vector<VarOrRVar> & vars);
Stage::predicate_stores(VarOrRVar v);
Stage::predicate_loads(VarOrRVar v);
Stage::split(VarOrRVar old, VarOrRVar outer, VarOrRVar inner, TailStrategy ts)
Stage::vectorize(VarOrRVar var, Expr factor, TailStrategy ts)
Stage::unroll(VarOrRVar var, Expr factor, TailStrategy ts)
Stage::tile(VarOrRVar x, VarOrRVar xo, VarOrRVar yo, VarOrRVar xi, VarOrRVar yi, TailStrategy tsx, TailStrategy tsy)
Stage::tile(VarOrRVar x, VarOrRVar xo, VarOrRVar yo, VarOrRVar xi, VarOrRVar yi, TailStrategy tsxy)
Stage::partition(VarOrRVar var, Partition p)

Tail Strategies

Two fundamental ones:

• RoundUp: align_extent(factor) & simple split with split.exact == true
• ShiftInward: as in main

And backward compatible ones:

• GuardWithIf == RoundUp + guard_with_if()
• Predicate == RoundUp + predicate_stores() + predicate_loads()
• PredicateLoads == RoundUp + predicate_loads()
• PredicateStores == RoundUp + predicate_stores()
• ShiftInwardsAndBlend == ShiftInward + predicate_loads()
• RoundUpAndBlend == RoundUp + predicate_stores()

There are now two new combinations possible: ShiftInwards + predicate_stores() or guard_with_if()
For large loop extents, the stores will never be out of bounds, and these will be functionally equivalent to ShiftInwards.
For small loop extents, shifting inward might go before the min, which will now be protected, and therefore possible.
However, for small loop extents, RoundUp + predicate_stores()/guard_with_if() is the same.
So while technically possible, they offer no additional value over the other combinations.

Reordering lowering of Partitioning and Vectorize/Unroll

Simply reordering those causes some failures, and probably a bunch of missed simplifications:
Running the test suite fails these:

	 99 - correctness_constraints (Failed)                  correctness
	153 - correctness_extract_concat_bits (Failed)          correctness
	182 - correctness_gpu_async_copy (Subprocess aborted)   correctness
	183 - correctness_gpu_async_copy_errors (Failed)        correctness
	263 - correctness_likely (Failed)                       correctness
	314 - correctness_predicated_store_load (Failed)        correctness
	429 - correctness_vectorize_nested (Failed)             correctness
	430 - correctness_vectorize_varying_allocation_size (Subprocess aborted) correctness
	431 - correctness_vectorized_assert (Subprocess aborted) correctness
	824 - python_correctness_serialization (Failed)         python python_correctness

So, (crazy?) idea, in Lower:

1. Unroll & vectorize loops
2. Loop partition & Record loop parititioning actions taken
3. Revert to before unroll & vectorize
4. Apply recorded loop partitioning actions: here we have the ability to mark tail handling loops as serial.
5. Vectorize & unroll loops

Partion Policies

Currently, in the main branch.

Loop partitioning is controlled in two ways in the main branch: there are policies attached to For nodes, and there
are likely's injected into index clamping and GuardWithIf.
Currently in main, the likely's are injected in some deterministic way during split() logic.
The user has no control over the injection of likely's in the split().
Having .partition_always() would do nothing if split() did not inject likely.

Proposed set of partition policies

• Partition::Never
• Partition::Always
• Partition::Auto

If additionally, we manage to reorder partitioning with unrolling/vectorizing, we can also define these:

• Partition::SerializeEdges: splits of the tail-case handling loop with a separate Serial for.
• Partition::UnrollEdges: splits of the tail-case handling loop with a separate Unrolled for.

Note that guard_with_if() does not have a Partition policy.
It does not make sense to schedule a guard_with_if() without a split.
Nothing would ever go out of bounds without a split.
Partition Policies are tighly bound together with split() or anything that calls split() internally.

Should we remove likely_if_innermost()?

As split() now fundamentally takes a Partition policy, we can immediately mark the corresponding for loop with that.
However, regarding the likely()s, there are two possible approaches:

1. let the partition policy influence the injection of likely() or likely_if_innermost().
2. always use likely() and have loop partitioning be smart when the policy is set to Auto.

Likely's and partition policy are two tools that are meant to control the same thing.

The first approach is to let Partion influence the injection of likely's.

• Partition::Never does not inject likely's
• Partition::Always does inject likely's
• Partition::Auto does inject likely_if_innermost

However, that leaves a problem with the boundary conditions:
If boundary conditions did not inject likely's unconditionally, but instead accept a Partition (which is API breaking), then that could controls the injection of likely's, in which case there is no reason to have parition() as a scheduling directive anymore.
However, that would not be able to communicate how we want to handle the edge-case: serial, unroll, or vectorized.

So, instead, likely() communicates what is the steady state; and the loop partitioning lowering pass should should look at the policy to decide what to do with that info.
This is a neat separation of concerns.

However, now we have this awkward likely_if_innermost(), which is currently used by ShiftInwards and Parition::Auto.
I find this to be a hack, and would like to remove likely_if_innermost() and use likely() everywhere where we expect steady state.
Instead, Loop Partitioning should decide for for-loops marked with Partition::Auto if the likely() it found to be good reason to partition or not.
I think the "smart" rule for Auto could be simple:

• match a Min(Ramp(loop_var)) or Max(Ramp(loop_var)) and you decide to partition the loop of loop_var (this typically happens in a vectorized x loop)

which implicitely then covers the case when you're in an index-clamped y loop:
Ramp(Min(loop_var)) or Ramp(Max(loop_var)) have the Ramp and Min/Max swapped inside-out.

This way boundary conditions are also handled.

One pitfall.

A user might do:

f.never_partition_all()
 .tile(x, y, xo, yo, xi, yi, 16, 16); // omits TailStrategy, and Partition!

As the .tile() call calls .split() internally, the API should probably NOT overwrite the partition policy which was already set by never_partition_all() back to Auto (due to Auto being the default argument when omitting it).
In that case, whenever a split() gets an Auto, we should not overwrite the existing policy.
Let's assume all loops are initialized with Partition::Auto.

What does this get us?

New scheduling possibility

You can now do: .align_bounds(x, 6).guard_with_if(x) on output buffers.
.align_bounds will align both the min and the extent, by increasing the total computed region of the Func.
.guard_with_if(x) will protect against out of bounds writes to the output buffer.
This is now a trivial results from combinatoric possibilities, as opposed to "making a version of align_bounds that predicates".

Note that .vectorize() and .unroll() both take a Partition policy, as they first call split() internally.
This gives us extremely clear scheduling on what happens:

// An innermost vectorized for-loop, where we recompute a bit of overlap in the last tile.
f.vectorize(x, TailStrategy::ShiftInward, Partition::Never);

// An innermost vectorized for-loop, with boundary condition needs to be split, to prevent a gather.
f.vectorize(x, TailStrategy::ShiftInward, Partition::Always);

// An innermost vectorized for-loop, but with simple for loop containing an if on the boundary conditions.
// This is basically a "RoundDown" vectorized loop + serial for loop for the tail .
f.vectorize(x, TailStrategy::GuardWithIf, Partition::SerializeEdges);

// An innermost vectorized for-loop, with RoundUp predicated vector stores/loads.
f.vectorize(x, TailStrategy::GuardWithIf, Partition::Never);

// An innermost vectorized for-loop non-predicated + a single predicated vector stores/loads inner iteratio.
f.vectorize(x, TailStrategy::GuardWithIf, Partition::Always);

// An innermost vectorized for-loop, "RoundDown" + unrolled loop-with-if
f.vectorize(x, TailStrategy::GuardWithIf, Partition::UnrollEdges);

Or, when demosaicing a 6x6 X-trans pattern:

f.align_bounds(x, 6, pattern_offset_x) // this implicitly rounds up
 .align_bounds(y, 6, pattern_offset_y) // this implicitly rounds up
 .guard_with_if(x, y) 
 .tile(x, y, xo, yo, xi, yi, 6, 6, TailStrategy::RoundUp, Parition::UnrollEdges)
	 // round up: only full tiles
	 // partition: create separate loop for edges and unroll it and replicate the injected if.
 .unroll(xi)
 .unroll(yi)

which is equivalent to:

f.align_bounds(x, 6, pattern_offset_x) // this implicitly rounds up
 .align_bounds(y, 6, pattern_offset_y) // this implicitly rounds up
 .tile(x, y, xo, yo, xi, yi, 6, 6, TailStrategy::GuardWithIf, Parition::UnrollEdges)
	 // guard with if: only full tiles, but guarded by ifs
	 // partition: create separate loop for edges and unroll it and replicate the injected if.
 .unroll(xi)
 .unroll(yi)

[abadams]

Don't guard_with_if/predicate_loads/predicate_stores also need to take a Partition to control if and how loops are partitioned to remove the guards?

[mcourteaux]

Ah, in case of the .align_bounds(x).guard_with_if(x) scenario you mean! I'd argue that you'd have to add .partition(x, your choice). No need to extra conflate that IMO. Thinking about this; I think the Partition argument can be removed everywhere if we take the approach of having likely() everywhere.