Blame view

stim/math/complexfield.cuh 3.44 KB
81e0d221   David Mayerich   separated executa...
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