如何在cupy的RawKernel中运行复杂的指数"cexpf"或"cexp"？

Question

作为标题，我正在计算由cupy提供的RawKernel中一个复数数组的指数。但我不知道如何包含或调用函数&;cexpf&;或";cexp"正确。错误信息总是显示";cexpf"是未定义的。有人知道如何以正确的方式调用这个函数吗?谢谢你的回答。

import cupy as cp
import time
add_kernel = cp.RawKernel(r'''
#include <cupy/complex.cuh>
#include <cupy/complex/cexpf.h>
extern "C" __global__
void test(double* x, double* y, complex<float>* z){
int tId_x = blockDim.x*blockIdx.x + threadIdx.x;
int tId_y = blockDim.y*blockIdx.y + threadIdx.y;

complex<float> value = complex<float>(x[tId_x],y[tId_y]);
z[tId_x*blockDim.y*gridDim.y+tId_y] = cexpf(value);
}''',"test")
x = cp.random.rand(1,8,4096,dtype = cp.float32)
#x = cp.arange(0,4096,dtype = cp.uint32)
y = cp.random.rand(1,8,4096,dtype = cp.float32)
#y = cp.arange(4096,8192,dtype = cp.uint32)
z = cp.zeros((4096,4096), dtype = cp.complex64)
t1 = time.time()
add_kernel((128,128),(32,32),(x,y,z))
print(time.time()-t1)
print(z)

Answer 1

查看标题，似乎您应该只调用exp，而不需要自己包含cupy/complex/cexpf.h，因为它已经通过cupy/complex.cuh隐式包含。

add_kernel = cp.RawKernel(r'''
#include <cupy/complex.cuh>
extern "C" __global__
void test(double* x, double* y, complex<float>* z){
int tId_x = blockDim.x*blockIdx.x + threadIdx.x;
int tId_y = blockDim.y*blockIdx.y + threadIdx.y;

complex<float> value = complex<float>(x[tId_x],y[tId_y]);
z[tId_x*blockDim.y*gridDim.y+tId_y] = exp(value);
}''',"test")

一般来说，Cupy的自定义内核c++复数API取自Thrust，因此您可以参考Thrust文档。请跳过使用thrust::命名空间。

反过来，Thrusts API在很大程度上尝试实现c++std::complexAPI，所以当Thrust文档不够深入时，查看c++标准库文档可能也会有所帮助。只是要小心，因为Thrust可能不会提供所有相同的保证来避免GPU上的性能问题。