Skip to content

Commit 2c00d59

Browse files
AsTonyshmentAsTonyshment
authored
Feature: Generalize element-wise vector operators for complex operands (#7791)
1 parent d88b719 commit 2c00d59

17 files changed

Lines changed: 878 additions & 701 deletions

File tree

‎source/source_base/kernels/cuda/math_kernel_op_vec.cu‎

Lines changed: 98 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,11 @@ __global__ void vector_mul_real_kernel(const int size,
3232
}
3333
}
3434

35-
template <typename T>
35+
template <typename T, typename Operand>
3636
__global__ void vector_mul_vector_kernel(const int size,
3737
T* result,
3838
const T* vector1,
39-
const typename GetTypeReal<T>::type* vector2,
39+
const Operand* vector2,
4040
const bool add)
4141
{
4242
int i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -66,11 +66,11 @@ __global__ void vector_div_constant_kernel(const int size,
6666
}
6767
}
6868

69-
template <typename T>
69+
template <typename T, typename Operand>
7070
__global__ void vector_div_vector_kernel(const int size,
7171
T* result,
7272
const T* vector1,
73-
const typename GetTypeReal<T>::type* vector2)
73+
const Operand* vector2)
7474
{
7575
int i = blockIdx.x * blockDim.x + threadIdx.x;
7676
if (i < size)
@@ -79,13 +79,13 @@ __global__ void vector_div_vector_kernel(const int size,
7979
}
8080
}
8181

82-
template <typename T, typename Real>
83-
__global__ void constantvector_addORsub_constantVector_kernel(const int size,
84-
T* result,
85-
const T* vector1,
86-
const Real constant1,
87-
const T* vector2,
88-
const Real constant2)
82+
template <typename T, typename Scalar>
83+
__global__ void vector_add_vector_kernel(const int size,
84+
T* result,
85+
const T* vector1,
86+
const Scalar constant1,
87+
const T* vector2,
88+
const Scalar constant2)
8989
{
9090
int i = blockIdx.x * blockDim.x + threadIdx.x;
9191
if (i < size)
@@ -190,123 +190,89 @@ void vector_div_constant_op<std::complex<double>, base_device::DEVICE_GPU>::oper
190190
vector_div_constant_wrapper(dim, result, vector, constant);
191191
}
192192

193-
// vector operator: result[i] = vector1[i](not complex) * vector2[i](not complex)
194-
template <>
195-
void vector_mul_vector_op<double, base_device::DEVICE_GPU>::operator()(const int& dim,
196-
double* result,
197-
const double* vector1,
198-
const double* vector2,
199-
const bool& add)
193+
template <typename T>
194+
inline T to_device_value(const T value)
200195
{
201-
int thread = thread_per_block;
202-
int block = (dim + thread - 1) / thread;
203-
vector_mul_vector_kernel<double><<<block, thread>>>(dim, result, vector1, vector2, add);
204-
205-
CHECK_CUDA_SYNC();
196+
return value;
206197
}
207-
// vector operator: result[i] = vector1[i](complex) * vector2[i](not complex)
208-
template <typename FPTYPE>
209-
inline void vector_mul_vector_complex_wrapper(const int& dim,
210-
std::complex<FPTYPE>* result,
211-
const std::complex<FPTYPE>* vector1,
212-
const FPTYPE* vector2,
213-
const bool& add)
214-
{
215-
thrust::complex<FPTYPE>* result_tmp = reinterpret_cast<thrust::complex<FPTYPE>*>(result);
216-
const thrust::complex<FPTYPE>* vector1_tmp = reinterpret_cast<const thrust::complex<FPTYPE>*>(vector1);
217-
int thread = thread_per_block;
218-
int block = (dim + thread - 1) / thread;
219-
vector_mul_vector_kernel<thrust::complex<FPTYPE>><<<block, thread>>>(dim, result_tmp, vector1_tmp, vector2, add);
220198

221-
CHECK_CUDA_SYNC();
222-
}
223-
template <>
224-
void vector_mul_vector_op<std::complex<float>, base_device::DEVICE_GPU>::operator()(const int& dim,
225-
std::complex<float>* result,
226-
const std::complex<float>* vector1,
227-
const float* vector2,
228-
const bool& add)
229-
{
230-
vector_mul_vector_complex_wrapper(dim, result, vector1, vector2, add);
231-
}
232-
template <>
233-
void vector_mul_vector_op<std::complex<double>, base_device::DEVICE_GPU>::operator()(
234-
const int& dim,
235-
std::complex<double>* result,
236-
const std::complex<double>* vector1,
237-
const double* vector2,
238-
const bool& add)
199+
template <typename FPTYPE>
200+
inline thrust::complex<FPTYPE> to_device_value(const std::complex<FPTYPE> value)
239201
{
240-
vector_mul_vector_complex_wrapper(dim, result, vector1, vector2, add);
202+
return thrust::complex<FPTYPE>(value.real(), value.imag());
241203
}
242204

243-
// vector operator: result[i] = vector1[i](not complex) / vector2[i](not complex)
244-
template <>
245-
void vector_div_vector_op<double, base_device::DEVICE_GPU>::operator()(const int& dim,
246-
double* result,
247-
const double* vector1,
248-
const double* vector2)
205+
template <typename T, typename Operand>
206+
void vector_mul_vector_op<T, base_device::DEVICE_GPU, Operand>::operator()(const int& dim,
207+
T* result,
208+
const T* vector1,
209+
const Operand* vector2,
210+
const bool& add)
249211
{
250-
int thread = thread_per_block;
251-
int block = (dim + thread - 1) / thread;
252-
vector_div_vector_kernel<double><<<block, thread>>>(dim, result, vector1, vector2);
212+
if (dim <= 0)
213+
{
214+
return;
215+
}
253216

217+
using DeviceT = typename GetTypeThrust<T>::type;
218+
using DeviceOperand = typename GetTypeThrust<Operand>::type;
219+
auto result_tmp = reinterpret_cast<DeviceT*>(result);
220+
auto vector1_tmp = reinterpret_cast<const DeviceT*>(vector1);
221+
auto vector2_tmp = reinterpret_cast<const DeviceOperand*>(vector2);
222+
const int thread = thread_per_block;
223+
const int block = (dim + thread - 1) / thread;
224+
vector_mul_vector_kernel<DeviceT, DeviceOperand>
225+
<<<block, thread>>>(dim, result_tmp, vector1_tmp, vector2_tmp, add);
254226
CHECK_CUDA_SYNC();
255227
}
256-
// vector operator: result[i] = vector1[i](complex) / vector2[i](not complex)
257-
template <typename FPTYPE>
258-
inline void vector_div_vector_complex_wrapper(const int& dim,
259-
std::complex<FPTYPE>* result,
260-
const std::complex<FPTYPE>* vector1,
261-
const FPTYPE* vector2)
228+
229+
template <typename T, typename Operand>
230+
void vector_div_vector_op<T, base_device::DEVICE_GPU, Operand>::operator()(const int& dim,
231+
T* result,
232+
const T* vector1,
233+
const Operand* vector2)
262234
{
263-
thrust::complex<FPTYPE>* result_tmp = reinterpret_cast<thrust::complex<FPTYPE>*>(result);
264-
const thrust::complex<FPTYPE>* vector1_tmp = reinterpret_cast<const thrust::complex<FPTYPE>*>(vector1);
265-
int thread = thread_per_block;
266-
int block = (dim + thread - 1) / thread;
267-
vector_div_vector_kernel<thrust::complex<FPTYPE>><<<block, thread>>>(dim, result_tmp, vector1_tmp, vector2);
235+
if (dim <= 0)
236+
{
237+
return;
238+
}
268239

240+
using DeviceT = typename GetTypeThrust<T>::type;
241+
using DeviceOperand = typename GetTypeThrust<Operand>::type;
242+
auto result_tmp = reinterpret_cast<DeviceT*>(result);
243+
auto vector1_tmp = reinterpret_cast<const DeviceT*>(vector1);
244+
auto vector2_tmp = reinterpret_cast<const DeviceOperand*>(vector2);
245+
const int thread = thread_per_block;
246+
const int block = (dim + thread - 1) / thread;
247+
vector_div_vector_kernel<DeviceT, DeviceOperand>
248+
<<<block, thread>>>(dim, result_tmp, vector1_tmp, vector2_tmp);
269249
CHECK_CUDA_SYNC();
270250
}
271-
template <>
272-
void vector_div_vector_op<std::complex<float>, base_device::DEVICE_GPU>::operator()(const int& dim,
273-
std::complex<float>* result,
274-
const std::complex<float>* vector1,
275-
const float* vector2)
276-
{
277-
vector_div_vector_complex_wrapper(dim, result, vector1, vector2);
278-
}
279-
template <>
280-
void vector_div_vector_op<std::complex<double>, base_device::DEVICE_GPU>::operator()(
281-
const int& dim,
282-
std::complex<double>* result,
283-
const std::complex<double>* vector1,
284-
const double* vector2)
285-
{
286-
vector_div_vector_complex_wrapper(dim, result, vector1, vector2);
287-
}
288251

289-
// vector operator: result[i] = vector1[i] * constant1 + vector2[i] * constant2
290-
template <typename T>
291-
void vector_add_vector_op<T, base_device::DEVICE_GPU>::operator()(const int& dim,
292-
T* result,
293-
const T* vector1,
294-
const Real constant1,
295-
const T* vector2,
296-
const Real constant2)
252+
template <typename T, typename Scalar>
253+
void vector_add_vector_op<T, base_device::DEVICE_GPU, Scalar>::operator()(const int& dim,
254+
T* result,
255+
const T* vector1,
256+
const Scalar constant1,
257+
const T* vector2,
258+
const Scalar constant2)
297259
{
298-
using Type = typename GetTypeThrust<T>::type;
299-
using Real = typename GetTypeReal<T>::type;
300-
301-
auto result_tmp = reinterpret_cast<Type*>(result);
302-
auto vector1_tmp = reinterpret_cast<const Type*>(vector1);
303-
auto vector2_tmp = reinterpret_cast<const Type*>(vector2);
304-
305-
int thread = thread_per_block;
306-
int block = (dim + thread - 1) / thread;
307-
constantvector_addORsub_constantVector_kernel<Type, Real>
308-
<<<block, thread>>>(dim, result_tmp, vector1_tmp, constant1, vector2_tmp, constant2);
260+
if (dim <= 0)
261+
{
262+
return;
263+
}
309264

265+
using DeviceT = typename GetTypeThrust<T>::type;
266+
using DeviceScalar = typename GetTypeThrust<Scalar>::type;
267+
auto result_tmp = reinterpret_cast<DeviceT*>(result);
268+
auto vector1_tmp = reinterpret_cast<const DeviceT*>(vector1);
269+
auto vector2_tmp = reinterpret_cast<const DeviceT*>(vector2);
270+
const DeviceScalar device_constant1 = to_device_value(constant1);
271+
const DeviceScalar device_constant2 = to_device_value(constant2);
272+
const int thread = thread_per_block;
273+
const int block = (dim + thread - 1) / thread;
274+
vector_add_vector_kernel<DeviceT, DeviceScalar>
275+
<<<block, thread>>>(dim, result_tmp, vector1_tmp, device_constant1, vector2_tmp, device_constant2);
310276
CHECK_CUDA_SYNC();
311277
}
312278

@@ -373,20 +339,28 @@ template struct vector_div_constant_op<std::complex<float>, base_device::DEVICE_
373339
template struct vector_div_constant_op<double, base_device::DEVICE_GPU>;
374340
template struct vector_div_constant_op<std::complex<double>, base_device::DEVICE_GPU>;
375341

376-
template struct vector_mul_vector_op<float, base_device::DEVICE_GPU>;
377-
template struct vector_mul_vector_op<std::complex<float>, base_device::DEVICE_GPU>;
378-
template struct vector_mul_vector_op<double, base_device::DEVICE_GPU>;
379-
template struct vector_mul_vector_op<std::complex<double>, base_device::DEVICE_GPU>;
380-
template struct vector_div_vector_op<std::complex<float>, base_device::DEVICE_GPU>;
381-
template struct vector_div_vector_op<double, base_device::DEVICE_GPU>;
382-
template struct vector_div_vector_op<std::complex<double>, base_device::DEVICE_GPU>;
383-
384-
template struct vector_add_vector_op<float, base_device::DEVICE_GPU>;
385-
template struct vector_add_vector_op<std::complex<float>, base_device::DEVICE_GPU>;
386-
template struct vector_add_vector_op<double, base_device::DEVICE_GPU>;
387-
template struct vector_add_vector_op<std::complex<double>, base_device::DEVICE_GPU>;
342+
template struct vector_mul_vector_op<float, base_device::DEVICE_GPU, float>;
343+
template struct vector_mul_vector_op<double, base_device::DEVICE_GPU, double>;
344+
template struct vector_mul_vector_op<std::complex<float>, base_device::DEVICE_GPU, float>;
345+
template struct vector_mul_vector_op<std::complex<double>, base_device::DEVICE_GPU, double>;
346+
template struct vector_mul_vector_op<std::complex<float>, base_device::DEVICE_GPU, std::complex<float>>;
347+
template struct vector_mul_vector_op<std::complex<double>, base_device::DEVICE_GPU, std::complex<double>>;
348+
349+
template struct vector_div_vector_op<float, base_device::DEVICE_GPU, float>;
350+
template struct vector_div_vector_op<double, base_device::DEVICE_GPU, double>;
351+
template struct vector_div_vector_op<std::complex<float>, base_device::DEVICE_GPU, float>;
352+
template struct vector_div_vector_op<std::complex<double>, base_device::DEVICE_GPU, double>;
353+
template struct vector_div_vector_op<std::complex<float>, base_device::DEVICE_GPU, std::complex<float>>;
354+
template struct vector_div_vector_op<std::complex<double>, base_device::DEVICE_GPU, std::complex<double>>;
355+
356+
template struct vector_add_vector_op<float, base_device::DEVICE_GPU, float>;
357+
template struct vector_add_vector_op<double, base_device::DEVICE_GPU, double>;
358+
template struct vector_add_vector_op<std::complex<float>, base_device::DEVICE_GPU, float>;
359+
template struct vector_add_vector_op<std::complex<double>, base_device::DEVICE_GPU, double>;
360+
template struct vector_add_vector_op<std::complex<float>, base_device::DEVICE_GPU, std::complex<float>>;
361+
template struct vector_add_vector_op<std::complex<double>, base_device::DEVICE_GPU, std::complex<double>>;
388362

389363
template struct dot_real_op<std::complex<float>, base_device::DEVICE_GPU>;
390364
template struct dot_real_op<double, base_device::DEVICE_GPU>;
391365
template struct dot_real_op<std::complex<double>, base_device::DEVICE_GPU>;
392-
} // namespace ModuleBase
366+
} // namespace ModuleBase

0 commit comments

Comments
 (0)