Show unflattened tensor in tvm.lower()

print(tvm.lower()) is very helpful. However, the tensor subscripts are flattened to 1D, which makes it a bit hard to read. Is there a way to keep the original dimension in tvm.lower() output?

Thanks.

Yuan

Well, find the answer myself (kind of).

Using the following

with tvm.build_config(dump_pass_ir=True) as cfg:
print(tvm.lower(s, [A, B, C], simple_mode=True))

The IR dumped in the first output file (0_ScheduleOps_ir.cc) has multi-dim subscripts.

2 Likes

Interesting, but this is pretty early in the IR passes.
So it is not really a representation of the final lowered version.

It is not.

For my use case, I’d like to see how the schedule affects the loop structure. Showing the final lowered version seems to obfuscate the code.

Hmmm I did something similar for VTA (as in wanted to see how the original loops get transformed) but to be honest I never remember seeing nonflattened arrays.

I see that in /python/tvm/build_module.py @364 stmt = ir_pass.StorageFlatten(stmt, binds, 64, cfg.instrument_bound_checkers)is the first pass of phase 1. So I would assume anything beyond phase 0 has 1D arrays.
Maybe there is no routine to undo this when printing. Maybe it could be added ?