[Unity][Transform] Implement relax.transform.ReorderTakeAfterMatmul

[Lunderberg]

If R.matmul(x, R.take(weights, indices)) occurs, with R.take selecting along the output feature dimension, it can be rearranged to R.take(R.matmul(x, weights), indices).

[tqchen]

what is the usecase of this pass? Seems such reordering would increase the amount of compute

[Lunderberg]

This pass would be part of an optimization pipeline that could be used for batched LoRA models. In this usage, the R.take(weights, indices) would select which LoRA should be used for each prompt in the batch. By rearranging the matmul, the resulting R.matmul(x, weights) could be combined with R.matmul(x, base_weights) using CombineParallelMatmul.

While this individual pass does increase the total amount of compute, the overall pipeline should (dependent on model size, number of loras, and batch size) improve performance by changing from three matmuls (one large matmul of the base weights and two small matmuls with the LoRA components) to two matmuls (one large matmul of concat(base_weights, lora_a) and one small matmul with lora_b).

[tqchen]

I feel for LoRA maybe the best approach we need some specialized kernels. cc @yzh119 to see if you have more suggestions based on your experience

[slyubomirsky]

Similar to the discussion in #16313, I don't think it would be harmful to include this pass as a tool for other transformations even if we don't plan to use it in the standard compilation immediately.

[tqchen]

i agree as long as it is optional

[Lunderberg]

Sounds good, and thank you! I should have led with this being an optional pass to be used as needed, not intended as part of a default pass.

[Lunderberg]

Rebased onto main, now that the main branch includes unity. No other changes made, just wanting to avoid stale CI before merge.