[TE] Add LegalizeInvalidAttach to legalize the compute_at location after split or fuse

[merrymercy]

Add LegalizeInvalidAttach in Schedule::normalize to legalize the compute_at location if the target iterator of compute_at is split or fused.

• Case 1: If the target of compute_at is split, we will move the compute_at location to the inner iterator.
• Case 2: If the target of compute_at is fused, we will move the compute_at location to the newly fused iterator.
  Case 2 can only happen if the target of compute_at is the innermost operand of fuse operation.

Examples: the following two cases will crash the compiler before this fix, but they are legal after this fix.

A=te.compute((10, 10), lambdai, j: 1.0, name='A')
B=te.compute((10, 10), lambdai, j: A[i][j], name='B')
# Case 1: Split an axis which is the target of a compute_ats=te.create_schedule([B.op])
s[A].compute_at(s[B], B.op.axis[1])
s[B].split(B.op.axis[1], 2)
print(tvm.lower(s, [A, B], simple_mode=True))
# Case 2: Fuse an axis which is the target of a compute_ats=te.create_schedule([B.op])
s[A].compute_at(s[B], B.op.axis[1])
s[B].fuse(B.op.axis[0], B.op.axis[1])
print(tvm.lower(s, [A, B], simple_mode=True))

[merrymercy]

This is a support PR for #5883. cc @tqchen@comaniac@jcf94

[comaniac]

Overall LGTM in terms of the functionality.

[merrymercy]

@tqchen@comaniac comments are addressed

[tqchen]

cc @vinx13@Hzfengsy@spectrometerHBH please also help to take a look, we can proceed to merge after reviews by a few more eyes

[vinx13]

LGTM

[junrushao]

Otherwise LGTM

[junrushao]

Maybe we don’t need this line

[merrymercy]

It is produced by my Jetbrain Clion