Add FFT support via AbstractFFTs interface

Project.toml

@@ -6,6 +6,7 @@ version = "1.9.3"
 projects = ["test", "docs", "perf", "examples", "res/wrap"]
 
 [deps]
+AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
@@ -25,6 +26,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Random123 = "74087812-796a-5b5d-8853-05524746bad3"
 RandomNumbers = "e6cf234a-135c-5ec9-84dd-332b85af5143"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
 ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -37,6 +39,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 SpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
+AbstractFFTs = "1"
 Adapt = "4.5"
 BFloat16s = "0.5, 0.6"
 CEnum = "0.4, 0.5"
@@ -56,6 +59,7 @@ Printf = "1"
 Random = "1"
 Random123 = "1.7.1"
 RandomNumbers = "1.6.0"
+Reexport = "1.2.2"
 SHA = "0.7, 1"
 ScopedValues = "1.3.0"
 SpecialFunctions = "2"

lib/mpsgraphs/MPSGraphs.jl

@@ -48,5 +48,6 @@ include("operations.jl")
 include("random.jl")
 
 include("matmul.jl")
+include("fft.jl")
 
 end

lib/mpsgraphs/fft.jl

@@ -0,0 +1,40 @@
+
+## FFT Descriptor Creation
+
+"""
+    MPSGraphFFTDescriptor(; inverse=false, scalingMode=MPSGraphFFTScalingModeNone)
+
+Create an MPSGraphFFTDescriptor with the specified parameters.
+"""
+function MPSGraphFFTDescriptor(; inverse::Bool = false, scalingMode::MPSGraphFFTScalingMode = MPSGraphFFTScalingModeNone)
+    obj = @objc [MPSGraphFFTDescriptor alloc]::id{MPSGraphFFTDescriptor}
+    desc = MPSGraphFFTDescriptor(obj)
+    desc.inverse = inverse
+    desc.scalingMode = scalingMode
+    return desc
+end
+
+## MPSGraph FFT operations
+function fastFourierTransformWithTensor(graph::MPSGraph, tensor::MPSGraphTensor, axes::NSArray, descriptor::MPSGraphFFTDescriptor, name = "fft")
+    obj = @objc [graph::id{MPSGraph} fastFourierTransformWithTensor:tensor::id{MPSGraphTensor}
+                                axes:axes::id{NSArray}
+                                descriptor:descriptor::id{MPSGraphFFTDescriptor}
+                                name:name::id{NSString}]::id{MPSGraphTensor}
+    MPSGraphTensor(obj)
+end
+
+function realToHermiteanFFTWithTensor(graph::MPSGraph, tensor::MPSGraphTensor, axes::NSArray, descriptor::MPSGraphFFTDescriptor, name = "rfft")
+    obj = @objc [graph::id{MPSGraph} realToHermiteanFFTWithTensor:tensor::id{MPSGraphTensor}
+                                axes:axes::id{NSArray}
+                                descriptor:descriptor::id{MPSGraphFFTDescriptor}
+                                name:name::id{NSString}]::id{MPSGraphTensor}
+    MPSGraphTensor(obj)
+end
+
+function HermiteanToRealFFTWithTensor(graph::MPSGraph, tensor::MPSGraphTensor, axes::NSArray, descriptor::MPSGraphFFTDescriptor, name = "irfft")
+    obj = @objc [graph::id{MPSGraph} HermiteanToRealFFTWithTensor:tensor::id{MPSGraphTensor}
+                                axes:axes::id{NSArray}
+                                descriptor:descriptor::id{MPSGraphFFTDescriptor}
+                                name:name::id{NSString}]::id{MPSGraphTensor}
+    MPSGraphTensor(obj)
+end

src/Metal.jl

@@ -14,6 +14,8 @@ import ObjectiveC: is_macos
 import KernelAbstractions
 using ScopedValues
 
+using Reexport: @reexport
+
 include("version.jl")
 include("storage_type.jl")
 
@@ -68,6 +70,7 @@ include("mapreduce.jl")
 include("accumulate.jl")
 include("indexing.jl")
 include("random.jl")
+include("fft.jl")
 
 # KernelAbstractions
 include("MetalKernels.jl")