[BACKEND] Add a loop unroller pass

[htyu]

Adding a loop unroller pass which applies to only loops with unroll annotation.

An annotated loop will look like:

    scf.for %arg5 = %c0_i32 to %arg3 step %c32_i32 : i32 {
      ...
    } {tt.loop_unroll_factor = 2 : i32}

[htyu]

#4639 is needed to unblock a test failure where the MLIR loop unroller doesn't work on for loops with integer IV.

[manman-ren]

This diff looks good to me!

[manman-ren]

Do you currently get performance wins from this without needing other changes?

[htyu]

Do you currently get performance wins from this without needing other changes?

Yes when used with kernel override, for kernels where register pressure isn't an issue (e.g persistent kernels). For E2E run we'll need a frontend change.

[ThomasRaoux]

I'm curious about the use case and if you considered using tl.static_range?

[htyu]

I'm curious about the use case and if you considered using tl.static_range?

One typical use case is that for the GEMM persistent kernel, like #4662 proposes, to annotate the flattened loop. An unroll factor of 2 on tile 128x128x128 gives some speedup, where SMEM and register pressure aren't an issue. The speedup can be higher when combined with our subsequent work of branch removal, i.e, removing useless branch checking for the prolog and epilog checks.

[ThomasRaoux]

I'm curious about the use case and if you considered using tl.static_range?

One typical use case is that for the GEMM persistent kernel, like #4662 proposes, to annotate the flattened loop. An unroll factor of 2 on tile 128x128x128 gives some speedup, where SMEM and register pressure aren't an issue. The speedup can be higher when combined with our subsequent work of branch removal, i.e, removing useless branch checking for the prolog and epilog checks.

but GEMM kernels usually don't have a static K loop? It doesn't sound like something that can be used in general cases?

[ThomasRaoux]

I'm curious about the use case and if you considered using tl.static_range?

One typical use case is that for the GEMM persistent kernel, like #4662 proposes, to annotate the flattened loop. An unroll factor of 2 on tile 128x128x128 gives some speedup, where SMEM and register pressure aren't an issue. The speedup can be higher when combined with our subsequent work of branch removal, i.e, removing useless branch checking for the prolog and epilog checks.

but GEMM kernels usually don't have a static K loop? It doesn't sound like something that can be used in general cases?

actually looking at the code again I see that it supports dynamic case. Is that what you do for GEMM? I don't understand how this helps since the unrolled iteration will have to be guarded by an IF op right?

[htyu]

I'm curious about the use case and if you considered using tl.static_range?

One typical use case is that for the GEMM persistent kernel, like #4662 proposes, to annotate the flattened loop. An unroll factor of 2 on tile 128x128x128 gives some speedup, where SMEM and register pressure aren't an issue. The speedup can be higher when combined with our subsequent work of branch removal, i.e, removing useless branch checking for the prolog and epilog checks.

but GEMM kernels usually don't have a static K loop? It doesn't sound like something that can be used in general cases?

actually looking at the code again I see that it supports dynamic case. Is that what you do for GEMM?

Yes, something like

`for _ in tl.range(0, k_tiles * tiles_per_SM, unroll_factor=F):`

where F comes from the autotuner. So far we've been giving it 2.

I don't understand how this helps since the unrolled iteration will have to be guarded by an IF op right?

So each original iteration comes with two if-checks, one to identify the start of the original conceptual inner K-loop, one is to identify the last iteration of the K-loop. Once unrolled by 2, there'd be four if-checks, two of which are unnecessary.

Yes the unrolled version comes with an epilog reminder loop, which can be unnecessary too. I think one more hint to the loop unroller should get it.

[ThomasRaoux]

I'm curious about the use case and if you considered using tl.static_range?

One typical use case is that for the GEMM persistent kernel, like #4662 proposes, to annotate the flattened loop. An unroll factor of 2 on tile 128x128x128 gives some speedup, where SMEM and register pressure aren't an issue. The speedup can be higher when combined with our subsequent work of branch removal, i.e, removing useless branch checking for the prolog and epilog checks.

but GEMM kernels usually don't have a static K loop? It doesn't sound like something that can be used in general cases?

actually looking at the code again I see that it supports dynamic case. Is that what you do for GEMM?

Yes, something like

`for _ in tl.range(0, k_tiles * tiles_per_SM, unroll_factor=F):`

where F comes from the autotuner. So far we've been giving it 2.

I don't understand how this helps since the unrolled iteration will have to be guarded by an IF op right?

So each original iteration comes with two if-checks, one to identify the start of the original conceptual inner K-loop, one is to identify the last iteration of the K-loop. Once unrolled by 2, there'd be four if-checks, two of which are unnecessary.

Yes the unrolled version comes with an epilog reminder loop, which can be unnecessary too. I think one more hint to the loop unroller should get it.

interesting, that makes sense

[htyu]

@ThomasRaoux @antiagainst Thanks for reviewing this change! How does the latest version look to you? Please let me know if you have more comments and I'm happy to address them.

[ThomasRaoux]

Looks fine to me assuming the FE changes are ready as well

[htyu]

Looks fine to me assuming the FE changes are ready as well

#4662 is for the FE side changes BTW.