merrymercy opened a new pull request #5917:
URL: https://github.com/apache/incubator-tvm/pull/5917
Add `LegalizeInvalidAttach` in `Schedule::normalize` to legalize the
compute_at location if the target iterator of compute_at was 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.
```python
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)
stmt = tvm.lower(s, [A, B], simple_mode=True)['main'].body
assert isinstance(stmt.body.body, tvm.tir.stmt.For)
# 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])
stmt = tvm.lower(s, [A, B], simple_mode=True)['main'].body
assert isinstance(stmt, tvm.tir.stmt.For)
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
