Allow output to alias intermediate tensor, step-2 to solve group norm segmentation issue #2375

[liqiangxl]

This PR is step-2 to solve #2375
Allow output to alias intermediate when:
(1) called from pre segment pass
(2) the op that preduces output interfer with reduction, see outputInterferingReduction

After this PR, the last reshape (which produces the ouput tensor) in group norm is changed to a no-op, see newly added test OutputAliasIntermediate

Util function getRepresentativeReductionTv is added for convenience and will be used to simipify reduction/normalization schedulers.

[liqiangxl]

!build

[wujingyue]

Thanks for the patch. Have you read #2395 (comment)? I tagged you there but the notification apparently didn't land. Would that solve this groupnorm issue?

[liqiangxl]

Yes and it looks great. IIUC, step-1 in the proposal is Run alias analysis as we do today, however, current AliasAnalysisResult::finalize needs to walk up the chain to find an input or output root and it will miss the case in group norm where the output alias an intermediate tensor. That's why I need to add this patch.
Previously, I also need to patch MarkAliasesPreparePass::runPass (see commit ), but with #2529, that patch is no longer required.

[wujingyue]

AliasAnalysisResult::finalize needs to walk up the chain to find an input or output root and it will miss the case in group norm where the output alias an intermediate tensor.

Yes, that's an important implementation detail that needs to be reworked. alias_to_source_ does keep aliases between intermediate tensors. So I'll have to put segment_set based on that without having to find the fusion I/O.

Overall, I feel lots of this PR is going to be superseded. However, if you need it in a week, I don't mind reviewing it. Wdyt?

[liqiangxl]

superseded

AliasAnalysisResult::finalize needs to walk up the chain to find an input or output root and it will miss the case in group norm where the output alias an intermediate tensor.

Yes, that's an important implementation detail that needs to be reworked. alias_to_source_ does keep aliases between intermediate tensors. So I'll have to put segment_set based on that without having to find the fusion I/O.

Overall, I feel lots of this PR is going to be superseded. However, if you need it in a week, I don't mind reviewing it. Wdyt?

Thanks. It would be great to have this PR merged, so we can make group norm works as expected. Assuming it won't cause troubles for your planned works to revise alias analysis.

[liqiangxl]

!build

[wujingyue]

I'd hope it's backed by numbers, but a more reasonable approach seems to be walk up until we reach a tensor that are known to be global, e.g., when outputInterferingReduction returns true or, not in this PR, the input of slice or pad.

[liqiangxl]

A little confusion about a tensor that are known to be global. Here I am trying to let the output tensor alias to a tensor that is not in global memory (not input or output, but intermediate), then further code will segment this output tensor = someOp(an intermediate tensor) to a no-op.

Here outputInterferingReduction depends on the whole fusion not a single tensor. So I didn't move it into the walk up process. But it looks like we should also check alias in alias_to_source_ is actually defined by a view op. Then we don't need to walk up. So the code is changed as:
(1) stop_at_view = may_alias_intermediate && outputInterferingReduction(fusion) is moved to finalize()
(2) only walk up if if (!isOpsToStop(alias->definition(), stop_at_view))

[wujingyue]

Makes sense. As discussed offline, it may make more sense to walk up until you find a non-view. E.g., in ... -> t0 -> (View) -> t1 -> (View) -> out, it's better to segment at t0 instead of t1.

[liqiangxl]

!build

[liqiangxl]

!build