c++ - 当数据在设备中时，有什么方法可以将 thrust 与 cufftComplex 数据类型一起使用？

Question

我正在制作具有 4D 数据的程序。在一个循环中，我遍历第 4 维，并获取最初为 float 的 3D 数据，我将该数据转换为 cufftComplex，然后将其复制到设备并执行多个设备傅里叶变换(来自 cufft 库)、图像去噪(我自己的 cuda 内核)等操作。在傅立叶变换之后处于这些操作之间的循环中时，我想在该特定迭代中找到 3D 数据的中值。

我知道 thrust 允许我进行排序，从中我可以很容易地找到中值，就像在下面的示例中一样，我使用普通的 float 数据类型

int row, column, slice;
row = 5;
column = 5;
slice = 5;
int alloc_size = row*column*slice;
float * theMainData = new float[alloc_size];
srand(time(NULL));
for(int k=0; k<slice; k++)
{
    for(int j=0; j<column; j++)
    {
       for(int i=0; i<row; i++)
       {
          theMainData[i+(j*column)+(k*row*column)] = rand()%1000;;
       }
    }
}
thrust::sort(theMainData, theMainData + alloc_size);
float median = theMainData[(alloc_size/2)];
cout<<endl<<"The median value is "<<median<<endl;

那么，现在我的问题是，当数据已经复制到设备时，有没有什么方法可以让 thrust::sort() 函数与 cufftComplex 一起工作？或者有没有类似thrust的函数可以在数据已经复制到设备的情况下处理cufftComplex数据类型？

Answer 1

So, now my question is, is there any way to make thrust::sort() function work with cufftComplex when the data already copied to the device?

是的。

Thrust提供 versions sort 操作接受用户定义的比较运算符，以您希望的方式对数据进行排序。

推力 sorting example 中给出了自定义仿函数的示例用法(使用 evens_before_odds 仿函数)。

所以你需要一个比较仿函数。我不知道你如何对复杂的值进行排序，但作为一个例子，如果你想根据你可以做的大小对它们进行排序:

#include <cuComplex.h>
#include <cufft.h>


struct my_complex_sort
{
  __host__ __device__
  bool operator()(cufftComplex x, cufftComplex y)
  {
    return (cuCabs(x) > cuCabs(y));
  }
};

如果设备上已有 cufftComplex 数据，则可以为其创建推力指针。假设指向您的设备数据的指针是 d_data 并且数据大小是 N:

thrust::device_ptr<cufftComplex> dp_data = thrust::device_pointer_cast(d_data);

然后您可以对数据进行排序:

thrust::sort(dp_data, dp_data+N, my_complex_sort());

(免责声明:在浏览器中编码，未经测试)