[TorchToLinalg] Add aten.fft_rfft and lowering

[giacs-epic]

• Add AtenFftRfftOp to Torch dialect.
• Add conversion of AtenFftRfftOp to Linalg, using a linalg.matmul per output component (real and imaginary). Computing the DFT is O(n^2).
• Add decomposition of AtenFftRfftOp into Torch-level ops (same paradigm as above).
• Add unit and end-to-end tests.

[zjgarvey]

Before reviewing in detail, let me see if I understand correctly what your cross-repository goal is.

1. You include this torch-to-linalg conversion as a baseline conversion, but in IREE, you intend to have a more performant lowering to linalg_ext? I suppose this pass exists before torch-to-linalg in the torch-to-iree pipeline?
2. Would it make more sense to add this as a decomposition of the op at the torch-dialect level? That way, other backends like Tosa and StableHLO could benefit, and we can turn off the op via backend-legal-ops option in torch-decompose-complex-ops pass if we want to go a different route in IREE. I have plans to modify the decompose complex ops pass to be more specific in the torch-to-iree pipeline this week, so we can specify a backend-legal-ops set there.

[giacs-epic]

@zjgarvey

1. Yes, that's the goal so far. Indeed I would place a pass to lower rfft to linalg_ext before torch-to-linalg in iree.
2. That makes a lot of sense. Would it be compatible with an eventual decomposition of aten.stft? I.e. would having both decompositions yield to the following behavior: aten.stft gets decomposed into aten.fft_rffts, which in turn get decomposed into matmuls?

[zjgarvey]

2. That makes a lot of sense. Would it be compatible with an eventual decomposition of `aten.stft`? I.e. would having both decompositions yield to the following behavior: `aten.stft` gets decomposed into `aten.fft_rfft`s, which in turn get decomposed into matmuls?

Yeah, precisely. There are some limitations, however. Does the higher performance path for fft_rfft to linalg_ext apply to the same cases as this conversion? If not, we will definitely need to keep this as a torch-to-linalg conversion to catch any patterns that failed to match the conversion to linalg_ext. This is because we won't be able to go back to decompose-complex-ops after trying to convert to linalg_ext.

[zjgarvey]

Ah, I see you already converted this to a decomposition. Perhaps we should just do both? StableHlo and Tosa would benefit from the decomposition, which we can turn off once you add the torch-to-linalg-ext path to IREE, and then the torch-to-linalg conversion would be a final fallback if the linalg_ext path doesn't apply.

[giacs-epic]

@zjgarvey The higher-performance path would apply when the input signal length is a power of 2, all other cases would need to be translated to this "naive" algorithm. Do you think it's possible to branch compilation based on the input dimension size?
Otherwise I'm open to just keeping both decomposition and lowering to linalg.

[zjgarvey]

It might be possible to mark the op as conditionally illegal for decompose-complex-ops, but I don't think we want to go that route. Let's add both the torch-to-linalg conversion and the decomposition for now.

[giacs-epic]

@zjgarvey Added conversion back.