Incorrect compute_at schedule can make TVM stuck. Below is the code:
import tvm
A = tvm.te.placeholder([6, 6], name="A")
C = tvm.te.compute([6, 6], lambda i, j: A[i, j] + 1, name="C")
sch = tvm.te.create_schedule(C.op)
print(tvm.lower(sch, [A, C], simple_mode=True))
sch[C].compute_at(sch[C], C.op.axis[0])
print("compute_at pass")
print(tvm.lower(sch, [A, C], simple_mode=True))
print("lower pass")
When running the above code, TVM will print “compute_at pass” but get stuck (forever) on “tvm.lower”, which is unexpected. I think that it could be better if some warning message or exception can be thrown to check whether the schedule is legal or not. Thanks!