Hilbert 变换#

1
#include<stdio.h>
2
#include<stdlib.h>
3
#include<string.h>
4
#include<cuda_runtime.h>
5
#include <cufft.h>
6
extern "C" {
7
    #include<rsf.h>
8
}
9

10
__global__ void apply_hilbert_phase_shift(cufftComplex* input, cufftComplex* output, int size) {
11
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
12
    if (idx < size) {
13
        // Apply phase shift for the Hilbert transform
14
        float phaseShift = (idx > size / 2) ? 1.0f : -1.0f;
15
        float tmp = input[idx].y;
16
        output[idx].y = input[idx].x*phaseShift;
17
        output[idx].x = -tmp*phaseShift;
18
    }
19
}
20

21
__global__ void scale_data(float *data, cufftComplex* output, int n, float scale)
22
{
23
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
24
    if (idx < n) {
25
       output[idx].x = data[idx] *scale;
26
       output[idx].y = 0.0f;
27
    }
28
}
29

30
__global__ void get_data(float *data, cufftComplex* output, int n)
31
{
32
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
33
    if (idx < n) {
34
        data[idx]= output[idx].x;
35
    }
36
}
37

38
void  hilbert_freq(float *indata, float *hilb,  int n1, int n2)
39
{
40
    cufftHandle plan;
41
    cufftComplex* inputComplex;
42
    cufftComplex* outputComplex;
43
    cufftResult result;
44
    cufftPlan1d(&plan, n1, CUFFT_C2C, 1);
45
    cudaMalloc((void**)&inputComplex, n1 * sizeof(cufftComplex));
46
    cudaMalloc((void**)&outputComplex, n1 * sizeof(cufftComplex));
47
    int threadsPerBlock = 256;
48
    int blocksPerGrid = (n1 + threadsPerBlock - 1) / threadsPerBlock;
49
    for(int i2=0; i2<n2; i2++){
50
        cudaMemset(outputComplex, 0, n1 * sizeof(cufftComplex));
51
        cudaMemset(inputComplex,  0, n1 * sizeof(cufftComplex));
52
        scale_data<<<(n1+255)/256, 256>>>(&indata[i2*n1], inputComplex, n1, 1.0f /n1);
53
        result = cufftExecC2C(plan, inputComplex, inputComplex,  CUFFT_FORWARD);
54
        if (result != CUFFT_SUCCESS) {
55
            fprintf(stderr, "CUFFT error: FFT execution failed, code %d \n", result);
56
            exit(1);
57
        }
58
        apply_hilbert_phase_shift<<<blocksPerGrid, threadsPerBlock>>>(inputComplex, outputComplex, n1);
59
        result = cufftExecC2C(plan,  outputComplex, outputComplex,  CUFFT_INVERSE);
60
        get_data<<<(n1+255)/256, 256>>>(&hilb[i2*n1], outputComplex, n1);
61
        if (result != CUFFT_SUCCESS) {
62
            fprintf(stderr, "CUFFT error: FFT execution failed, code %d \n", result);
63
            exit(1);
64
        }
65
    }
66
    cufftDestroy(plan);
67
    cudaFree(inputComplex);
68
    cudaFree(outputComplex);
69
}
70

71
int main(int argc, char *argv[]){
72

73
    int n1,n2;
74
    float *data, *hilb;
75
    float *d_data, *d_hilb;
76
    sf_file in, out;
77

78
    sf_init (argc,argv);
79

80
    in = sf_input("in");
81
    out = sf_output("out");
82

83
    if (SF_FLOAT != sf_gettype(in)) sf_error("Need float input");
84
    if (!sf_histint(in,"n1",&n1)) sf_error("No n1= in input");
85
    n2 = sf_leftsize(in,1);
86

87
    data = sf_floatalloc(n1*n2);
88
    hilb = sf_floatalloc(n1*n2);
89

90
    cudaMalloc((float**)&d_data,n1*n2*sizeof(float)); cudaMemset(d_data,0,n1*n2*sizeof(float));
91
    cudaMalloc((float**)&d_hilb,n1*n2*sizeof(float)); cudaMemset(d_hilb,0,n1*n2*sizeof(float));
92

93
    sf_floatread(data,n1*n2,in); cudaMemcpy(d_data,data,n1*n2*sizeof(float),cudaMemcpyHostToDevice);
94

95
    hilbert_freq(d_data, d_hilb,  n1, n2);
96

97
    cudaMemcpy(hilb, d_hilb, n1*n2*sizeof(float), cudaMemcpyDeviceToHost);
98

99
    sf_floatwrite(hilb,n1*n2,out);
100

101
    cudaFree(d_data);
102
    cudaFree(d_hilb);
103
    exit(0);
104
}