推力中的 sort_by_key 是阻塞调用吗？

Question

我反复将一系列内核排入队列：

for 1..100:
    for 1..10000:
        // Enqueue GPU kernels
        Kernel 1 - update each element of array  
        Kernel 2 - sort array  
        Kernel 3 - operate on array  
    end
    // run some CPU code
    output "Waiting for GPU to finish"
    // copy from device to host
    cudaMemcpy ... D2H(array)
end

内核 3 的顺序为 O(N^2)，因此是迄今为止最慢的。对于内核 2，我直接在设备上使用推力::sort_by_key：

thrust::device_ptr<unsigned int> key(dKey);
thrust::device_ptr<unsigned int> value(dValue);
thrust::sort_by_key(key,key+N,value);

似乎这个对推力的调用是阻塞的，因为 CPU 代码只有在内循环完成后才会执行。我看到这一点是因为如果我删除对sort_by_key 的调用，主机代码（正确地）在内循环完成之前输出“等待”字符串，而如果我运行排序则不会。

有没有办法异步调用thrust::sort_by_key？

Answer 1

1. 首先考虑有一个内核启动队列，它只能容纳这么多挂起的启动。一旦启动队列已满，任何类型的其他内核启动都会阻塞。在空队列槽可用之前，主机线程将不会继续（超出那些启动请求）。我很确定 3 次内核启动的 10000 次迭代将在达到 10000 次迭代之前填满这个队列。因此，如果您按顺序启动 30000 个内核，那么任何类型的非平凡内核启动都会有一些延迟（我认为）。 （然而，最终，当所有内核因为一些已经完成而被添加到队列中时，如果没有其他阻塞行为，那么您会在所有内核实际完成之前看到“等待...”消息。）

2. thrust::sort_by_keyrequires temporary storage（大小大约等于您的数据集大小）。每次您使用它时，都会通过 cudaMalloc 操作在后台分配此临时存储空间。这个cudaMalloc 操作是阻塞的。当cudaMalloc 从主机线程启动时，它会等待内核活动的间隙才能继续。

要解决第 2 项，似乎至少有两种可能的方法：

1. 提供thrust custom allocator。根据此分配器的特性，您可能能够消除阻塞 cudaMalloc 行为。 （但请参阅下面的讨论）

2. 使用cub SortPairs。这里的优点（如我所见 - 您的示例不完整）是您可以进行一次分配（假设您知道整个循环迭代中最坏情况的临时存储大小）并消除在您的内部进行临时内存分配的需要环形。

据我所知，推力方法（上面的 1）仍然会在每次迭代时有效地执行某种临时分配/释放步骤，即使您提供自定义分配器也是如此。如果您有一个设计良好的自定义分配器，那么这可能几乎是一个“无操作”。 cub 方法似乎有需要知道最大大小的缺点（为了完全消除分配/释放步骤的需要），但我认为对于推力自定义分配器也有同样的要求。否则，如果您需要在某个时候分配更多内存，自定义分配器实际上将不得不执行类似cudaMalloc 之类的操作，这将在工作中发挥作用。