CUDA排序Z轴3D阵列C++/Thrust

我正在寻找排序一个大型三维数组沿着z轴。
示例数组为X x Y x Z（1000x1000x5）
我想沿着z轴排序，所以我会沿着z轴对5个元素执行1000x1000次排序。
编辑更新：尝试使用下面的thrust。它是可用的，我会把输出存储回去，但是这非常慢，因为我在每个（x，y）位置一次排序5个元素：

#include <stdio.h>
#include <stdlib.h>
#include <iostream>

#include <thrust/device_ptr.h>
#include <thrust/sort.h>
#include <thrust/gather.h>
#include <thrust/iterator/counting_iterator.h>

int main(){
int x = 1000, y = 1000, z = 5;
float*** unsorted_cube = new float** [x];

for (int i = 0; i < x; i++) 
{
    // Allocate memory blocks for 
    // rows of each 2D array 
    unsorted_cube[i] = new float* [y];

    for (int j = 0; j < y; j++) 
    {
        // Allocate memory blocks for 
        // columns of each 2D array 
        unsorted_cube[i][j] = new float[z];
    }
}

for (int i = 0; i < x; i++)
{
    for (int j = 0; j < y; j++)
    {
        unsorted_cube[i][j][0] = 4.0f;
        unsorted_cube[i][j][1] = 3.0f;
        unsorted_cube[i][j][2] = 1.0f;
        unsorted_cube[i][j][3] = 5.0f;
        unsorted_cube[i][j][4] = 2.0f;
    }
}

for (int i = 0; i < 5; i++)
{
    printf("unsorted_cube first 5 elements to sort at (0,0): %f\n", unsorted_cube[0][0][i]);
}

float* temp_input;
float* temp_output;
float* raw_ptr;
float raw_ptr_out[5];
cudaMalloc((void**)&raw_ptr, N_Size * sizeof(float));
for (int i = 0; i < x; i++)
{ 
    for (int j = 0; j < y; j++)
    {
        temp_input[0] = unsorted_cube[i][j][0];
        temp_input[1] = unsorted_cube[i][j][1];
        temp_input[2] = unsorted_cube[i][j][2];
        temp_input[3] = unsorted_cube[i][j][3];
        temp_input[4] = unsorted_cube[i][j][4];

        cudaMemcpy(raw_ptr, temp_input, 5 * sizeof(float), cudaMemcpyHostToDevice);
        thrust::device_ptr<float> dev_ptr = thrust::device_pointer_cast(raw_ptr);
        thrust::sort(dev_ptr, dev_ptr + 5);
        thrust::host_vector<float> host_vec(5);
        thrust::copy(dev_ptr, dev_ptr + 5, raw_ptr_out);

        if (i == 0 && j == 0)
        {
            for (int i = 0; i < 5; i++)
            {
                temp_output[i] = raw_ptr_out[i];
            }
            printf("sorted_cube[0,0,0] : %f\n", temp_output[0]);
            printf("sorted_cube[0,0,1] : %f\n", temp_output[1]);
            printf("sorted_cube[0,0,2] : %f\n", temp_output[2]);
            printf("sorted_cube[0,0,3] : %f\n", temp_output[3]);
            printf("sorted_cube[0,0,4] : %f\n", temp_output[4]);
        }
    }
}
}

假设数据的格式为每个xy平面中的值在存储器中连续：data[((z * y_length) + y) * x_length + x]（也最适合合并GPU上的内存访问）

#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/for_each.h>
#include <thrust/zip_iterator.h>

void sort_in_z_dir(thrust::device_vector<float> &data,
                   int x_length, int y_length) { // z_length == 5
  auto z_stride = x_length * y_length;
  
  thrust::for_each(
      thrust::make_zip_iterator(thrust::make_tuple(
          data.begin(),
          data.begin() + z_stride,
          data.begin() + 2 * z_stride,
          data.begin() + 3 * z_stride,
          data.begin() + 4 * z_stride)),
      thrust::make_zip_iterator(thrust::make_tuple(
          data.begin() + z_stride,
          data.begin() + 2 * z_stride,
          data.begin() + 3 * z_stride,
          data.begin() + 4 * z_stride,
          data.begin() + 5 * z_stride)),
      [] __host__ __device__ 
      (thrust::tuple<float, float, float, float, float> &values) {
        float local_data[5] = {thrust::get<0>(values),
                               thrust::get<1>(values),
                               thrust::get<2>(values),
                               thrust::get<3>(values),
                               thrust::get<4>(values)};
        thrust::sort(thrust::seq, local_data, local_data + 5);
        thrust::get<0>(values) = local_data[0];
        thrust::get<1>(values) = local_data[1];
        thrust::get<2>(values) = local_data[2];
        thrust::get<3>(values) = local_data[3];
        thrust::get<4>(values) = local_data[4];
      });
}

这个解决方案在硬编码z_length方面确实非常丑陋。人们可以使用一些C++模板-“魔法”将z_length变成一个模板参数，但对于这个关于Thrust的答案来说，这似乎是矫枉过正。
有关数组和元组之间接口的示例，请参见Convert std::tuple to std::array C++11和How to convert std::array to std::tuple?。
这个解决方案的好处是排序算法本身应该是性能最优的，我不知道thrust::sort是否针对这么小的输入数组进行了优化，但是你可以用我在评论中提出的任何自己编写的排序算法来代替它。
如果你希望能够使用不同的z_length而不遇到这些麻烦，你可能更喜欢这个解决方案，它在全局内存中排序，这远非最优，而且感觉有点笨拙，因为它使用Thrust几乎只是为了启动内核。data[((x * y_length) + y) * z_length + z]

#include <thrust/counting_iterator.h>
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/for_each.h>

void sort_in_z_dir_alternative(thrust::device_vector<float> &data,
                               int x_length, int y_length, int z_length) {
  int n_threads = x_length * y_length;
  
  thrust::for_each(
      thrust::make_counting_iterator(0),
      thrust::make_counting_iterator(n_threads),
      [ddata = thrust::raw_pointer_cast(data.data()), z_length] __host__ __device__ (int idx) {
        thrust::sort(thrust::seq, 
                     ddata + z_length * idx,
                     ddata + z_length * (idx + 1));
      });
}

如果您同意将z_length作为模板参数，那么这可能是一个结合了两个领域优点的解决方案（数据格式与第一个示例类似）：

#include <thrust/counting_iterator.h>
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/for_each.h>

template <int z_length>
void sort_in_z_dir_middle_ground(thrust::device_vector<float> &data,
                                 int x_length, int y_length) {
  int n_threads = x_length * y_length; // == z_stride
  
  thrust::for_each(
      thrust::make_counting_iterator(0),
      thrust::make_counting_iterator(n_threads),
      [ddata = thrust::raw_pointer_cast(data.data()),
       z_length, n_threads] __host__ __device__ (int idx) {
        float local_data[z_length];
        #pragma unroll
        for (int i = 0; i < z_length; ++i) {
          local_data[i] = ddata[idx + i * n_threads];
        }
        thrust::sort(thrust::seq, 
                     local_data,
                     local_data + z_length);
        #pragma unroll
        for (int i = 0; i < z_length; ++i) {
          ddata[idx + i * n_threads] = local_data[i];
        }
      });
}

CUDA排序Z轴3D阵列C++/Thrust

1条答案

相关问题

热门标签

最新问答