Show unflattened tensor in tvm.lower()

#1

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

#2

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
Print out schedule for debugging
[DISCUSS] Contributing new docs for InferBound
#3

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

#4

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.

#5

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 ?