查找金属纹理中的最小值和最大值

Posted 2023-02-15

【中文标题】查找金属纹理中的最小值和最大值

【英文标题】：Finding the minimum and maximum value within a Metal texture

【发布时间】：2016-08-08 09:36:37

【问题描述】：

我有一个包含 16 位无符号整数 (MTLPixelFormatR16Uint) 的 MTLTexture。值的范围从大约 7000 到 20000，其中 0 用作“nodata”值，这就是在下面的代码中跳过它的原因。我想找到最小值和最大值，以便我可以在 0-255 之间重新调整这些值。最终，我将寻找基于数据直方图的最小值和最大值（它有一些异常值），但现在我坚持简单地提取最小值/最大值。

我可以将数据从 GPU 读取到 CPU 并提取最小值/最大值，但我更愿意在 GPU 上执行此任务。

第一次尝试

每个线程组使用 16x16 个线程调度命令编码器，线程组的数量取决于纹理大小（例如，宽度 = textureWidth / 16，高度 = textureHeight / 16）。

typedef struct 
    atomic_uint min;
    atomic_uint max;
 BandMinMax;

kernel void minMax(texture2d<ushort, access::read> band1 [[texture(0)]],
                   device BandMinMax &out [[buffer(0)]],
                   uint2 gid [[thread_position_in_grid]])

    ushort value = band1.read(gid).r;

    if (value != 0) 
        uint currentMin = atomic_load_explicit(&out.min, memory_order_relaxed);
        uint currentMax = atomic_load_explicit(&out.max, memory_order_relaxed);

        if (value > currentMax) 
            atomic_store_explicit(&out.max, value, memory_order_relaxed);
        
        if (value < currentMin) 
            atomic_store_explicit(&out.min, value, memory_order_relaxed);
        
    

由此我得到一个最小值和最大值，但对于同一个数据集，最小值和最大值通常会返回不同的值。可以肯定的是，当有多个线程在运行时，这是单个线程的最小值和最大值。

第二次尝试

在上一次尝试的基础上，这次我将存储每个线程的单个最小值/最大值，全部为 256 (16x16)。

kernel void minMax(texture2d<ushort, access::read> band1 [[texture(0)]],
                   device BandMinMax *out [[buffer(0)]],
                   uint2 gid [[thread_position_in_grid]],
                   uint tid [[ thread_index_in_threadgroup ]])


    ushort value = band1.read(gid).r;

    if (value != 0) 
        uint currentMin = atomic_load_explicit(&out[tid].min, memory_order_relaxed);
        uint currentMax = atomic_load_explicit(&out[tid].max, memory_order_relaxed);

        if (value > currentMax) 
            atomic_store_explicit(&out[tid].max, value, memory_order_relaxed);
        
        if (value < currentMin) 
            atomic_store_explicit(&out[tid].min, value, memory_order_relaxed);
        
    

这将返回一个包含 256 组最小/最大值的数组。从这些我想我可以找到最小值中的最低值，但这似乎是一个糟糕的方法。将不胜感激指向正确方向的指针，谢谢！

【参考方案1】：

Metal Shading Language 具有原子比较和交换函数，可用于将内存位置的现有值与某个值进行比较，如果它们比较不相等，则替换该位置的值。有了这些，您可以创建一组原子比较和替换如果-[大于|小于]-比操作：

static void atomic_uint_exchange_if_less_than(volatile device atomic_uint *current, uint candidate)

    uint val;
    do 
        val = *((device uint *)current);
     while ((candidate < val || val == 0) && !atomic_compare_exchange_weak_explicit(current,
                                                                                     &val,
                                                                                     candidate,
                                                                                     memory_order_relaxed,
                                                                                     memory_order_relaxed));


static void atomic_uint_exchange_if_greater_than(volatile device atomic_uint *current, uint candidate)

    uint val;
    do 
        val = *((device uint *)current);
     while (candidate > val && !atomic_compare_exchange_weak_explicit(current,
                                                                       &val,
                                                                       candidate,
                                                                       memory_order_relaxed,
                                                                       memory_order_relaxed));

要应用这些，您可以创建一个缓冲区，其中每个线程组包含一个交错的最小、最大对。然后，在核函数中，从纹理中读取并有条件地写入最小值和最大值：

kernel void min_max_per_threadgroup(texture2d<ushort, access::read> texture [[texture(0)]],
                                    device uint *mapBuffer [[buffer(0)]],
                                    uint2 tpig [[thread_position_in_grid]],
                                    uint2 tgpig [[threadgroup_position_in_grid]],
                                    uint2 tgpg [[threadgroups_per_grid]])

    ushort val = texture.read(tpig).r;

    device atomic_uint *atomicBuffer = (device atomic_uint *)mapBuffer;

    atomic_uint_exchange_if_less_than(atomicBuffer + ((tgpig[1] * tgpg[0] + tgpig[0]) * 2),
                                      val);

    atomic_uint_exchange_if_greater_than(atomicBuffer + ((tgpig[1] * tgpg[0] + tgpig[0]) * 2) + 1,
                                         val);

最后，运行一个单独的内核来减少这个缓冲区并收集整个纹理的最终最小值、最大值：

kernel void min_max_reduce(constant uint *mapBuffer [[buffer(0)]],
                           device uint *reduceBuffer [[buffer(1)]],
                           uint2 tpig [[thread_position_in_grid]])

    uint minv = mapBuffer[tpig[0] * 2];
    uint maxv = mapBuffer[tpig[0] * 2 + 1];

    device atomic_uint *atomicBuffer = (device atomic_uint *)reduceBuffer;

    atomic_uint_exchange_if_less_than(atomicBuffer, minv);

    atomic_uint_exchange_if_greater_than(atomicBuffer + 1, maxv);

当然，您只能减少设备允许的总线程执行宽度（~256），因此您可能需要多次减少，每次减少要操作的数据大小最大线程执行宽度的一个因子。

免责声明：这可能不是最好的技术，但在我对 OS X 实现的有限测试中似乎是正确的。它比 Intel Iris Pro 上 256x256 纹理上的简单 CPU 实现略快，但在 Nvidia GT 750M 上慢得多（因为调度开销）。

【讨论】：

感谢@warrenm，似乎工作正常。我对原子缓冲区的偏移量有疑问；例如atomicBuffer + ((tgpig[1] * tpt[0] + tgpig[0]) * 2)。我的理解是原子操作适用于每个线程组（如果顺便说一句，请更正这些假设中的任何一个）？我通过 threads_per_threadgroup 注释传递到内核的线程组使用 16x16 线程到 tpt 变量。我不确定这是我的线程组网格的宽度吗？例如;纹理大小为 192x160，线程组网格为 12x10，偏移量计算为 atomicBuffer + ((tgpig[1] * 12 + tgpig[0]) * 2)?

请原谅最后一行硬编码的 12。我想我想说的是在 min_max_per_threadgroup 内核中用threadgroups_per_grid 替换threads_per_threadgroup 修复它？

@lock 是的，你完全正确。我的实现很幸运，因为threads_per_threadgroup 恰好等于threadgroups_per_grid。以上已更正。

TBH，这可能比一开始就对我的理解更有帮助。再次感谢。

Warrenm，我很好奇，你是如何将 Metal 代码执行从 Iris 切换到 Nvidia 的？