tvm icon indicating copy to clipboard operation
tvm copied to clipboard

[TIR] Update symbolic index term order in loop fusion

Open wrongtest-intellif opened this issue 1 month ago • 1 comments

This change just keep stride terms order the same with fused loop order in fuse primitive. In symbolic circumstances, previous form suffer from simplification issues and would make the expression tree much complex in following lowering steps.

Take [M, N] tiling as an example, the previous binding form after

i, j = sch.get_loops(block_b)
i0, i1 = sch.split(i, factors=[None, 64])
j0, j1 = sch.split(j, factors=[None, 16])
sch.reorder(i0, j0, i1, j1)
sch.fuse(i0, j0)

would be like (i_0_j_0_fused in [0, ceildiv(M, 64) * ceildiv(N, 16)]

vi = T.axis.spatial(M, i_0_j_0_fused % ((N + 15) // 16 * ((M + 63) // 64)) // ((N + 15) // 16) * 64 + i_1)

instead of more simple version

vi = T.axis.spatial(M, i_0_j_0_fused // ((N + 15) // 16) * 64 + i_1)

This is because unfortunately we do not know ceildiv(N, 16) * ceildiv(M, 64) == ceildiv(M, 64) * ceildiv(N, 16) in rule based simplifications. And then certain analysis (for example, region estimation) may fail to give concise estimations, due to complex dynamic expression trees.

wrongtest-intellif avatar Oct 30 '25 05:10 wrongtest-intellif