Make Flux Transformer RoPE use a custom IREE kernel

[sogartar]

We assume that the custom kernel would yield better performance instead of using PyTorch ops.

[sogartar]

To match the required dimension order, axis permutation is performed.
I am not sure if this can be optimized out. Needs some investigation of
the compilation passes.

This change introduces a significant (10X) deterioration in performance.
Baseline performance is 552 ms.
With this change it is 6131 ms.

[sogartar]

The large performance problem has been addressed by iree-org/iree#19822.

[renxida]

Nit: would it make more sense to just rename _compute_rotary_embedding_table?

[sogartar]

I want to use this function outside of the class.

[renxida]

Yes, i mean, instead of:

• copying compute_rotary_embedding_table so we can use it outside the class
• make old _compute_rotary_embedding_table redirect to compute_rotary_embedding_table

Just

• rename _compute_rotary_embedding_table to compute_rotary_embedding_table and use it outside the class
• change all referenes to _compute_rotary_embedding_table to use compute_rotary_embedding_table instead

The latter requires an IDE and is slightly more work, but does not leave a stub/redirect function behind.

[sogartar]

The class variant got updated to handle the hugging face case so it has more hair now.

[sogartar]

This PR is waiting on iree-org/iree#19829.