有没有一种方法可以在Thrust的一个内核调用中完成reduce_by_key
操作和reduce
(或者理想情况下是另一个reduce_by_key
)操作?除了获得计算速度之外,让我们假设我想这样做,因为第一个reduce_by_key
操作的输出值数量太大,无法存储在内存中。
我一直在想transform_output_iterator
是否可以在这里提供帮助,但还没有找到解决方案。
一个简单的演示(但不是我的真实的用例)可以是找到矩阵中每行最大值的最小值,其中该矩阵被展平并存储在device_vector
中。
1条答案
按热度按时间brqmpdu11#
下面的代码计算所有行最大值中的最小值,并使用固定数量的临时存储来存储有限数量的最小值。然后,执行minreduce以找到全局最小值
这个想法是通过transform_output_iterator直接更新最小值。这可以通过原子(在用于临时最小值的原始指针的情况下)或通过锁(在用于临时最小值的迭代器的情况下)来完成。未在本答案中显示)。
为了避免原子争用,临时最小值的数目不应太小。
对于大小为1的1G段,即每个输入元素都有一个原子操作,我在A100 GPU上观察到以下时序。