tvm
tvm copied to clipboard
[TIR] Update symbolic index term order in loop fusion
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.