[Triton] Make tl.cat deterministic by using permute+reshape+join and remove CatOp

[Mogball]

tl.cat is now implemented with permute, reshape, join, reshape, permute instead of using the builtin CatOp. This makes tl.cat deterministic, and thanks to linear layouts, it doesn't appear to affect performance much. The can_reorder option is deprecated and now does nothing.

tl.cat also supports concatenating along any dim.

[ThomasRaoux]

Yay, this is awesome!

[lezcano]

very cool!

[songdejun]

With Triton 3.3.1+git96394dff, when implementing the randn op, tl.pair_uniform_to_normal generates 2 * BLOCK random normal numbers, but in an interleaved memory layout. With tl.cat (the older implementation), these numbers could be reordered to be contiguous, and we didn't care whether CatOp would break determinism. This method brought huge performance improvements.

However, with Triton 3.5.0, the improvements disappear.

Is this the reason for the disappearance?

[lezcano]

@ThomasRaoux I think that we shoud have a tl.can_reorder for people to be able to do this wherever they want, rather than have it baked into our view ops now that our views are actual views.

[ThomasRaoux]

With Triton 3.3.1+git96394dff, when implementing the randn op, tl.pair_uniform_to_normal generates 2 * BLOCK random normal numbers, but in an interleaved memory layout. With tl.cat (the older implementation), these numbers could be reordered to be contiguous, and we didn't care whether CatOp would break determinism. This method brought huge performance improvements.

However, with Triton 3.5.0, the improvements disappear.

Is this the reason for the disappearance?

This change is not part of 3.5 so I don't think this is related.

@ThomasRaoux I think that we shoud have a tl.can_reorder for people to be able to do this wherever they want, rather than have it baked into our view ops now that our views are actual views.

Interesting idea, one thing that have been really annoying in the past is that if the value has multiple uses, they may want different layouts, in this case it may or may not be correct to use different layouts for different uses. It feels like the cases where being able to re-order elements in a reasonable way is rare, I think we should only enable it/keep it enabled if we find a strong use case.

For random value generation I think we should always be able to propagate layout to be able to get the output layout we want for free.

[lezcano]

Even better if we don't find a use case for reorder, but if there is, we could always just allow it asserting that it only has one user.

[ThomasRaoux]

we could always just allow it asserting that it only has one user.

It can be subtle though, as the user may get rematerialized by the compiler. But we can figure it out if there is a real life use case. For now except for the case of reshape + two step reduction I haven't seen a legit use case