Add support for FFT plans

[giordano]

Should fix #1926, but for some reason I don't quite understand, the in-place plans aren't in-place (the added tests should fail), even though the code in #1926 (comment) works as expected. I now get

julia> using Reactant, FFTW

julia> x_host = rand(ComplexF32, 16, 16); x_r = Reactant.to_rarray(x_host);

julia> planned_fft!(x) = plan_fft!(x) * x
planned_fft! (generic function with 1 method)

julia> @code_hlo planned_fft!(x_r)

module @reactant_planned... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<16x16xcomplex<f32>> {enzymexla.memory_effects = []}) -> tensor<16x16xcomplex<f32>> attributes {enzymexla.memory_effects = []} {
    %0 = stablehlo.fft %arg0, type =  FFT, length = [16, 16] : (tensor<16x16xcomplex<f32>>) -> tensor<16x16xcomplex<f32>>
    return %0 : tensor<16x16xcomplex<f32>>
  }
}

Note that contrary to #1926 (comment), the input is not aliasing the output. @avik-pal do you have any clue of what's going on here?

Edit: I simply missed a ! 🤦

[giordano]

I'd hope XLA is able to reuse plans applied to different arrays of the same shape. Since all the planning is completely opaque to the frontend we can't do much about it, but reusing plans is quite a crucial optimisation (and the whole point plans exist)

[wsmoses]

I'd assume so as well, but also since we already separate compile from run, during compile even if we make a different plan for all the fft instructions, they'd still all be planned during runtime ? It already does this with matmul by trying all possible matmul plans I know

[giordano]

The requisites for a plan are the shape and the eltype of the target array, and XLA should known them all at compile-time, right? So I'd speculate it's possible it does all the planning during compilation. But still, planning can be quite expensive, not repeating it unnecessarily would be welcome in large applications, whether it's during compile- or run-time.

[wsmoses]

yeah Id assume it has a cache, but we can benchmark, check, and fix, as needed as a follow up

[wsmoses]

whats up with rfft?

[giordano]

Rfft returns an array of different shape than the input one, it can't be done in-place. There's no plan_rfft!

[wsmoses]

ah got it!