complexfield.cuh
3.44 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
#ifndef RTS_COMPLEXFIELD_H
#define RTS_COMPLEXFIELD_H
#include "cublas_v2.h"
#include <cuda_runtime.h>
#include "../math/field.cuh"
#include "../math/complex.h"
#include "../math/realfield.cuh"
namespace stim{
template<typename T>
__global__ void gpu_complexfield_mag(T* dest, complex<T>* source, unsigned int r0, unsigned int r1){
int iu = blockIdx.x * blockDim.x + threadIdx.x;
int iv = blockIdx.y * blockDim.y + threadIdx.y;
//make sure that the thread indices are in-bounds
if(iu >= r0 || iv >= r1) return;
//compute the index into the field
int i = iv*r0 + iu;
//calculate and store the result
dest[i] = source[i].abs();
}
/*This class stores functions for saving images of complex fields
*/
template<typename T, unsigned int D = 1>
class complexfield : public field< stim::complex<T>, D >{
using field< stim::complex<T>, D >::R;
using field< stim::complex<T>, D >::X;
using field< stim::complex<T>, D >::shape;
using field< stim::complex<T>, D >::cuda_params;
public:
//find the maximum value of component n
stim::complex<T> find_max(unsigned int n){
cublasStatus_t stat;
cublasHandle_t handle;
//create a CUBLAS handle
stat = cublasCreate(&handle);
if(stat != CUBLAS_STATUS_SUCCESS){
std::cout<<"CUBLAS Error: initialization failed"<<std::endl;
exit(1);
}
int L = R[0] * R[1]; //compute the number of discrete points in a slice
int index; //result of the max operation
stim::complex<T> result;
if(sizeof(T) == 8)
stat = cublasIcamax(handle, L, (const cuComplex*)X[n], 1, &index);
else
stat = cublasIzamax(handle, L, (const cuDoubleComplex*)X[n], 1, &index);
index -= 1; //adjust for 1-based indexing
//if there was a GPU error, terminate
if(stat != CUBLAS_STATUS_SUCCESS){
std::cout<<"CUBLAS Error: failure finding maximum value."<<std::endl;
exit(1);
}
//retrieve the maximum value for this slice and store it in the maxVal array
std::cout<<X[n]<<std::endl;
HANDLE_ERROR(cudaMemcpy(&result, X[n] + index, sizeof(stim::complex<T>), cudaMemcpyDeviceToHost));
return result;
}
public:
enum attribute {magnitude, real, imaginary};
//constructor (no parameters)
complexfield() : field<stim::complex<T>, D>(){};
//constructor (resolution specified)
complexfield(unsigned int r0, unsigned int r1) : field<stim::complex<T>, D>(r0, r1){};
//assignment from a field of complex values
complexfield & operator=(const field< stim::complex<T>, D > rhs){
field< complex<T>, D >::operator=(rhs);
return *this;
}
//assignment operator (scalar value)
complexfield & operator= (const complex<T> rhs){
field< complex<T>, D >::operator=(rhs);
return *this;
}
//assignment operator (vector value)
complexfield & operator= (const vec< complex<T>, D > rhs){
field< complex<T>, D >::operator=(rhs);
return *this;
}
//cropping
complexfield crop(unsigned int width, unsigned int height){
complexfield<T, D> result;
result = field< complex<T>, D>::crop(width, height);
return result;
}
void toImage(std::string filename, attribute type = magnitude, unsigned int n=0){
field<T, 1> rf(R[0], R[1]);
//get cuda parameters
dim3 blocks, grids;
cuda_params(grids, blocks);
if(type == magnitude){
gpu_complexfield_mag <<<grids, blocks>>> (rf.ptr(), X[n], R[0], R[1]);
rf.toImage(filename, n, true);
}
}
};
} //end namespace rts
#endif