132 lines
2.4 KiB
Go
132 lines
2.4 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"math"
|
|
"math/cmplx"
|
|
|
|
"github.com/klauspost/cpuid/v2"
|
|
)
|
|
|
|
// FFT performs Fast Fourier Transform on complex data
|
|
func FFT(data []complex128) []complex128 {
|
|
if len(data) == 0 {
|
|
return data
|
|
}
|
|
|
|
// Check if we can use AVX512
|
|
if cpuid.CPU.AVX512F() && cpuid.CPU.AVX512DQ() {
|
|
return fftAVX512(data)
|
|
}
|
|
|
|
// Fallback to standard Go implementation
|
|
return fftGo(data)
|
|
}
|
|
|
|
// fftGo is the standard Go implementation of FFT
|
|
func fftGo(data []complex128) []complex128 {
|
|
n := len(data)
|
|
if n == 1 {
|
|
return data
|
|
}
|
|
|
|
// Ensure n is a power of 2
|
|
if n&(n-1) != 0 {
|
|
// Pad with zeros to next power of 2
|
|
nextPower := 1
|
|
for nextPower < n {
|
|
nextPower <<= 1
|
|
}
|
|
padded := make([]complex128, nextPower)
|
|
copy(padded, data)
|
|
data = padded
|
|
n = nextPower
|
|
}
|
|
|
|
// Bit-reversal permutation
|
|
rev := make([]int, n)
|
|
for i := 0; i < n; i++ {
|
|
rev[i] = rev[i>>1]>>1 | (i&1)<<int(math.Log2(float64(n))-1)
|
|
}
|
|
|
|
// Apply bit-reversal
|
|
result := make([]complex128, n)
|
|
for i := 0; i < n; i++ {
|
|
result[i] = data[rev[i]]
|
|
}
|
|
|
|
// Cooley-Tukey FFT
|
|
for size := 2; size <= n; size <<= 1 {
|
|
half := size >> 1
|
|
angle := -2 * math.Pi / float64(size)
|
|
w := complex(math.Cos(angle), math.Sin(angle))
|
|
|
|
for i := 0; i < n; i += size {
|
|
wi := complex(1, 0)
|
|
for j := 0; j < half; j++ {
|
|
t := wi * result[i+j+half]
|
|
result[i+j+half] = result[i+j] - t
|
|
result[i+j] += t
|
|
wi *= w
|
|
}
|
|
}
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
// fftAVX512 calls the AVX512 assembly implementation
|
|
//go:noescape
|
|
func fftAVX512(data []complex128) []complex128
|
|
|
|
// Inverse FFT
|
|
func IFFT(data []complex128) []complex128 {
|
|
n := len(data)
|
|
if n == 0 {
|
|
return data
|
|
}
|
|
|
|
// Conjugate input
|
|
conj := make([]complex128, n)
|
|
for i := 0; i < n; i++ {
|
|
conj[i] = cmplx.Conj(data[i])
|
|
}
|
|
|
|
// Apply FFT
|
|
fftResult := FFT(conj)
|
|
|
|
// Conjugate output and scale
|
|
result := make([]complex128, n)
|
|
for i := 0; i < n; i++ {
|
|
result[i] = cmplx.Conj(fftResult[i]) / complex(float64(n), 0)
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
func main() {
|
|
// Example usage
|
|
fmt.Println("AVX512 Support:", cpuid.CPU.AVX512F() && cpuid.CPU.AVX512DQ())
|
|
|
|
// Test data
|
|
data := []complex128{
|
|
complex(1, 0),
|
|
complex(2, 0),
|
|
complex(3, 0),
|
|
complex(4, 0),
|
|
complex(5, 0),
|
|
complex(6, 0),
|
|
complex(7, 0),
|
|
complex(8, 0),
|
|
}
|
|
|
|
fmt.Println("Input:", data)
|
|
|
|
// Forward FFT
|
|
fftResult := FFT(data)
|
|
fmt.Println("FFT Result:", fftResult)
|
|
|
|
// Inverse FFT
|
|
ifftResult := IFFT(fftResult)
|
|
fmt.Println("IFFT Result:", ifftResult)
|
|
} |