[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), lambda i, j: 1.0, name='A')
    B = te.compute((10, 10), lambda i, j: A[i][j], name='B')

    # Case 1: Split an axis which is the target of a compute_at
    s = 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_at
    s = 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