[TIR] Add structural error printing for TensorIR

[shingjan]

This PR intends to improve the error rendering by annotating regions of interest in TIR like synr following PR #8121.
Previously when there is an error:

The IR is:
@tvm.script.ir_module
class Module:
    @tir.prim_func
    def main(a: tir.handle, b: tir.handle) -> None:
        A = tir.match_buffer(a, [128, 128, 128, 128], dtype="float32")
        B = tir.match_buffer(b, [128, 128, 128, 128], dtype="float32")
        # body
        # with tir.block("root")
        for i, j, k, l in tir.grid(128, 128, 128, 8):
            with tir.block([128, 128, 128, 128], "B") as [vi, vj, vk, vl]:
                tir.bind(vi, i)
                tir.bind(vj, j)
                tir.bind(vk, k)
                tir.bind(vl, l * 16)
                tir.reads([A[vi, vj, vk, vl]])
                tir.writes([B[vi, vj, vk, vl]])
                B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * tir.float32(2)
    
Regions of interest:
tir.Block#0
block B(iter_var(vi, range(min=0, ext=128)), iter_var(vj, range(min=0, ext=128)), iter_var(vk, range(min=0, ext=128)), iter_var(vl, range(min=0, ext=128))) {
  reads([A[vi, vj, vk, vl]])
  writes([B[vi, vj, vk, vl]])
  B[vi, vj, vk, vl] = (A[vi, vj, vk, vl]*2f)
}

With this PR, the error will be rendered like below:

The IR with diagnostic is:
@tvm.script.ir_module
class Module:
    @tir.prim_func
    def main(a: tir.handle, b: tir.handle) -> None:
        A = tir.match_buffer(a, [128, 128, 128, 128], dtype="float32")
        B = tir.match_buffer(b, [128, 128, 128, 128], dtype="float32")
        # body
        # with tir.block("root")
        for i, j, k, l in tir.grid(128, 128, 128, 8):
            tir.Block#0
            with tir.block("B"):
            ^^^^^^^^^^^^^^^^^^^^
                vi, vj, vk = tir.axis.remap("SSS", [i, j, k])
                vl = tir.axis.spatial(128, l * 16)
                tir.reads([A[vi, vj, vk, vl]])
                tir.writes([B[vi, vj, vk, vl]])
                B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * tir.float32(2)

cc: @vinx13 @junrushao1994

[vinx13]

Thanks for the PR. I left some comments

[junrushao]

Hey thanks for the PR! It's a pretty nice POC, and I would like to deliberate on the format :-)

• Shall we update the syntax in the description to reflect the latest change on mainline? Also please add a few unittests
• If the error location is at a particular loop, let's avoid fusing it with other loops with tir.grid
• Also we should consider printing the names of the error locations under the annotations, so that the error message could be clearly referring to these blocks/loops

[junrushao]

Hmm I just went through the code, but figured that the snippet in the description isn't the actual output of the PR (Sorry I was wrong)? Shall we update the description with a real-world example?

We have plenty of these errors in tests/python/unittest/test_tir_schedule_*.py, guarded by:

with pytest.raises(tvm.tir.ScheduleError, ...):`
  ...

[vinx13]

left some comments

[Hzfengsy]

Looks good to me. One thing I'm not sure is that will it compatible with the new block syntax.

[tqchen]

@shingjan you may need to update the test cases and code now that the new block syntax is introduced

[shingjan]

Printing nested loop is fixed as well as some comments addressed. Should be good for another good before we merge this in. @vinx13 @junrushao1994

[vinx13]

otherwise LGTM

[vinx13]

LGTM, @Hzfengsy would you like to take a second look?

[Hzfengsy]

LGTM