pointwise scheduler fails to validate reference tv

[jjsjann123]

Fixes: #3512

When picking reference tv, pointwise scheduler fails to validate that the transformation on reference tv can be safely propagated to all outputs in the fusion. The issue occurs when an IterDomain that's not in the reference tv is merged with another dimension in the output tv, preventing the merge on reference tv to be propagated to the target.

This PR adds an optional check areAllOutputIdsMappedTo in nvfuser::pointwise_utils::DomainMap::isValidReference

The added check in this PR checks that all source producer IterDomain producing the IterDomain on outputs are covered by reference tv. This is safe for pointwise scheduler, since the scheduler checks that there's no reversible view present in the fusion.

The check is optional and is disabled by transpose scheduler, where the reference_tv is not supposed to cover the entire fusion, but rather a subset of fusion IO tensors. We should extent that in future PRs.

[jjsjann123]

!test

[jjsjann123]

!test

[jjsjann123]

🤞

[jjsjann123]

!test

[jjsjann123]

!test --diff-bench

[naoyam]

But now I'm worried that this check might not work for pad, where we go from a broadcast ID to non-broadcast ID... I'm checking that with a test.

Ah, this actually can cause an issue. For example, suppose we pick a tensor as a reference that has a broadcast ID, and that broadcast ID comes from a fusion input tensor. Suppose that broadcast ID is also used by a pad op, generating a non-broadcast ID, and that non-broadcast ID is NOT included in the reference.

More specifically:

t0 = [i0, b1] // fusion input
t1 = [i2, i3, i4, b5] // fusion input

// t1: [b6, b7, i0, b1]
t2 = broadcast(t0, [true, true, false, false]) 

t3 = add(t1, t2) // fusion output


// t4: [i0, i8]
t4 = pad(t0) // fusion output

Here, suppose we choose t3 as a reference. Clearly, we can't schedule the i8 domain of t4. However, since the propagation back from i8 reaches at b1, which is also reachable from t3, the reference validation would miss flagging the invalid reference.

[jjsjann123]

Clearly, we can't schedule the i8 domain of t4.

I was worried about the same thing and I was thinking about changing the get_source_iter_domains here to add output of resize to also be treated as sources.

But turns out we don't need that.
Transform propagation can actually schedule fusion like that from t3.

You can look at the other example I added for pad (there's a typo in the comment on tv0, I'll fix that). It's very similar to yours
I also just tried you example and the scheduling seems to work. I'll add it to the test as well.

I think the difference between this example and the original issue we had is due to resize. Which links the input ID to the output ID, so propagation of transformation works (I think?!) I was planning to ask you that question tomorrow, since it's getting late when I ran into this.

[jjsjann123]

!test --diff-bench

[jjsjann123]

!test --diff-bench

[jjsjann123]

errr... what's with CI 😭

[jjsjann123]

!test --diff-bench

[jjsjann123]

My gut feeling, based on my naive understanding about transform propagation, if we actually have an expandOp, or even just use a resize to link the IDs that are expanded, I feel nvfuser might be able to codegen the original issue #3512 without a problem.

The other scenario is also pretty interesting to me. #3576

NOTE for myself. I think I should go have a look at how transform propagation actually works to verify this for a piece of mind.

[naoyam]

LGTM. Thanks for the fix.

[jjsjann123]

!test --diff-bench

[jjsjann123]

!test --diff-bench