Translate the repetition pattern to expand and reshape

[naoyam]

In RoPE, this repeat pattern shows up commonly:

https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L136

torch.cat((freqs, freqs), dim=-1)

This pattern shows up as PadOp followed by CatOp in nvFuser.

This preseg pass translates this pattern to expand and reshape ops. For example, given a pattern like:

t0 = [i0];
t1 = cat({t0, t0}, -1);

It will be translated to:

t0 = [i0]
t2 = broadcast(t0, {true, false});
t3 = expand(t2, {2, i0});
t4 = reshape(t3, {2 * i0});

And all uses of t1 will be replaced by t4.

While the pattern can be handled by the resize scheduler, it's currently limited to segments with pointwise ops only, and its scheduling heuristics are not tuned yet. Specifically, I experimentally observed a significant perf gain with the Mistral backward function.

[naoyam]

!test --diff

[naoyam]

!test --diff

[naoyam]

!test --diff