FFT support

[AshtonSBradley]

FFT support is absent from Metal.jl right now.

A state of the art and fully feature fft library has been developed for Vulcan that has a Metal backed.

https://github.com/DTolm/VkFFT

Would it be possible to wrap it from Metal.jl?

[maleadt]

Sure, that's very possible. It's not high up my list of priorities though, so feel free to take a stab at it. It would require packaging vkFFT on Yggdrasil, probably using separate packages for each of the back-ends, and using that from here in Metal.jl.

[max-Hawkins]

This would be awesome!

[rveltz]

It seems they wrapped FFT in torch

pytorch/pytorch#119670

In [1]: import torch
   ...: x = torch.randn(1, 16000, device="mps")

In [2]: x
Out[2]: x

tensor([[-0.2197, -1.0679,  0.9474,  ..., -0.9471,  0.4943,  0.0663]],
       device='mps:0')

In [3]: x
Out[3]: 
tensor([[-0.2197, -1.0679,  0.9474,  ..., -0.9471,  0.4943,  0.0663]],
       device='mps:0')

In [4]: torch.fft.rfft(x)
Out[4]: 
tensor([[-153.5669+0.0000e+00j, -122.9657+7.1824e+01j,   19.3239-4.1716e+01j,
          ...,  -62.0741-5.7114e+01j,   37.2413-6.5813e+01j,
          -76.1166+1.1871e-05j]], device='mps:0')

[rveltz]

https://twitter.com/rohanpaul_ai/status/1793417620915626217

[maleadt]

Alternatively, we may be able to use MLX, which has a C API: ml-explore/mlx#1102

[christiangnrd]

This will be fairly straightforward to add once MPSGraph support is in (#566).

[rveltz]

now that it is merged, how far are we from fft?

[christiangnrd]

That depends on how long it takes for someone with the bandwidth and will to implement it to come along.

[bjarthur]

a jll exists for MLX: https://github.com/JuliaBinaryWrappers/MLX_jll.jl. am i correct in thinking then that this just needs to be added as a dep to Metal.jl and then some ccalls added somewhere?