Combining Separate tir.Block using compute_at()

Questions
10

Thanks @xiacijie for the question!

# Snippet-1:
for i, j in tir.grid(128, 128):
    with tir.block([128, 128], "init") as [vi, vj]:
        C[vi, vj] = 0.0

# Snippet-2:
with tir.block([128, 128], "init") as [vi, vj]:
    C[vi, vj] = 0.0

Yes, the two snippets above are strictly equivalent, and the second one is the syntactic sugar for the first. TVM script has a built-in functionality called “auto-completion” that desugars snippets.

To get a sense what the fully desugared IR looks like, you may print it out with the following command in python:

print(tvm.script.asscript(PrimFunc-or-IRModule))