Make non-divisible splits not change extents used in indexing

[naoyam]

This is a second attempt to fix an indexing problem that is first reported in PR #1243. See that PR as well for the context of this problem.

Also, the repro of issue #1074 is now working (FusionNonDivisibleSplit2).

This PR addresses an index linearlization problem that was exposed by view and strided gather operations, both of which are based on rFactor. The root cause of the problem is non-divisible splits. When indexing with such splits, the domain that is split gets a larger extent due to the ceilDiv of the outer domain. This is propagated back from leaf to parent domains in IndexCompute::handle(Split*). It turns out that there are consistency problems when to propagate. The most sensible way is likely to completely stop doing so, but it also means intermediate domains have to be predicated if split by a non-divisible factor.

Also, predicating such domains doesn't work if vectorization is involved. If an output domain of a split is vectorized, the split must be divisible. Divisibility is validated at the kernel launch time.

The main logic is in NonDivisibleSplitInfo. It scans a fusion to locate any Split expression that must be either predicated or validated. If an IterDomain is run-time validated, it's not predicated as it's unnecessary (NonDivisibleSplitInfo::removeRedundancy()).

getReferenceRootPredicates is extended to generate predicates for those domains that may be split with a non-divisible factor. Currently, it finds root or merged contiguous domains to predicate. I extended that to include non-divisible split domains (getNonDivisibleDomainsToPredicate). While index_compute.cpp may appear to have many changes, most of them don't change the existing logic. For example, the change in getStartAndStopOffsets is not to do adjustments for shift and gather operations when predicating non-divisible split domains.

All of the C++ tests are passing, but test_native_layer_norm_half is failing due to a runtime validation of a non-divisible domain with vectorization. The domain to split is:

i515 = T0.size[1] * T0.size[2]
~ i518 = T0.size[0] * i515
~ i521 = ceilDiv(i518, 4)
i521

This i521 domain is split by 2 and the inner domain of size 2 is vectorized. The test fails with T0 being [17, 10, 10] since 17 * 10 * 10 / 4 isn't divisible by 2. Without the runtime check, an alignment error occurs. Since @csarofeen mentioned he's working on changing scheduling algorithms for vectorization, I've done nothing on the failure yet.

Closes #1074

[naoyam]

This was a test created for #1250 but removed as it was failing due to the extent problem.

[csarofeen]

Should we clean up the predication logic to always use an explicit list of which iteration domains need to be predicated? I understand why this was only done for non-divisible-splits, but it seems to me we could have a pass that would specify all the domains that need to be predicated, including if root domains need predication, or just the contiguous merge domains. Maybe it would be good to try and unify this logic a bit more.

Looks really good overall, approving!

[csarofeen]

Can you mention in these assert messages this is from vectorizedSplit validation. A little more context in these messages will be helpful.

[csarofeen]

Nit: Temporarily

[csarofeen]

Nit: divisibility

[csarofeen]

Should non-const references in this header be private? Maybe we should check this as a follow up.

[naoyam]

As long as GpuLower is passed as a const object after lowering, I think it should be fine, but yes, this class is growing and may need some refactoring.

[csarofeen]

Nit: support shift. If a contig_id is a merged non-root domain, nothing

[naoyam]

Should we clean up the predication logic to always use an explicit list of which iteration domains need to be predicated? I understand why this was only done for non-divisible-splits, but it seems to me we could have a pass that would specify all the domains that need to be predicated, including if root domains need predication, or just the contiguous merge domains. Maybe it would be good to try and unify this logic a bit more.

Agreed. We have contig IDs, non-divisible split IDs, and non-exact parallelized IDs. The last one is added in getInlinePredicate, but there's no reason that all three can't be done uniformly. I'll revisit later.