fourier.cu

#include <stdlib.h>
#include <stdio.h>
#include <WolframLibrary.h>

#include <cuda_runtime.h>
#include <cufft.h>

//#include <helper_cuda.h>

typedef float2 Complex;
typedef double2 doubleComplex;

extern "C"
DLLEXPORT mint WolframLibrary_getVersion( ) {
        return WolframLibraryVersion;
}

extern "C"
DLLEXPORT int WolframLibrary_initialize( WolframLibraryData libData) {
        return 0;
}

extern "C"
DLLEXPORT void WolframLibrary_uninitialize( WolframLibraryData libData) {
        return;
}

extern "C"
DLLEXPORT int fourier(WolframLibraryData libData, mint argc, MArgument *args, MArgument res)
{
    MTensor signal  = MArgument_getMTensor(args[0]);
    int N = libData->MTensor_getDimensions(signal)[0];
    int mem_size = N*sizeof(doubleComplex);

    doubleComplex *h_signal = (doubleComplex*) libData->MTensor_getComplexData(signal);

    doubleComplex *d_signal;
    cudaMalloc((void **)&d_signal, mem_size);
    cudaMemcpy(d_signal, h_signal, mem_size, cudaMemcpyHostToDevice);

    cufftHandle plan;
    (cufftPlan1d(&plan, N, CUFFT_Z2Z, 1));

    cufftExecZ2Z(plan, (cufftDoubleComplex *)d_signal, (cufftDoubleComplex *)d_signal, CUFFT_INVERSE);

    cudaMemcpy(h_signal, d_signal, mem_size, cudaMemcpyDeviceToHost);

    cufftDestroy(plan);
    cudaFree(d_signal);

    MArgument_setMTensor(res, signal);
    return 0;
}

extern "C"
DLLEXPORT int fourier_single(WolframLibraryData libData, mint argc, MArgument *args, MArgument res)
{
    MTensor signal  = MArgument_getMTensor(args[0]);
    int N = libData->MTensor_getDimensions(signal)[0];
    int mem_size = N*sizeof(Complex);

    doubleComplex *h_signal_double = (doubleComplex*) libData->MTensor_getComplexData(signal);

    Complex *h_signal = (Complex*) malloc(mem_size);
    for (int i=0; i<N; ++i)
    {
        h_signal[i].x=h_signal_double[i].x;
        h_signal[i].y=h_signal_double[i].y;
    }

    doubleComplex *d_signal;
    cudaMalloc((void **)&d_signal, mem_size);
    cudaMemcpy(d_signal, h_signal, mem_size, cudaMemcpyHostToDevice);

    cufftHandle plan;
    cufftPlan1d(&plan, N, CUFFT_C2C, 1);

    cufftExecC2C(plan, (cufftComplex *)d_signal, (cufftComplex *)d_signal, CUFFT_INVERSE);

    cudaMemcpy(h_signal, d_signal, mem_size, cudaMemcpyDeviceToHost);

    for (int i=0; i<N; ++i)
    {
        h_signal_double[i].x=h_signal[i].x;
        h_signal_double[i].y=h_signal[i].y;
    }

    cufftDestroy(plan);
    cudaFree(d_signal);

    MArgument_setMTensor(res, signal);
    return 0;
}