查找Metal纹理中的最小值和最大值

Question

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

我可以将GPU中的数据读取到CPU并将最小/最大值拉出,但更愿意在GPU上执行此任务.

第一次尝试

命令编码器被调度为每个线程组16x16个线程,线程组的数量基于纹理大小(例如,width = textureWidth/16,height = 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);
        }
    }
}

从中我得到最小值和最大值,但对于相同的数据集,min和max通常会返回不同的值.相当确定这是多线程运行时来自单个线程的最小值和最大值.

第二次尝试

在前一次尝试的基础上,这次我将存储每个线程的各个最小/最大值,全部为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组最小/最大值的数组.从这些我想我可以找到最低值的最低值,但这似乎是一个糟糕的方法.希望指向正确的方向,谢谢!

Answer 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实现的有限测试中看起来确实是正确的.它比英特尔Iris Pro上的256x256纹理上的天真CPU实现略快,但在Nvidia GT 750M上却大大减慢(因为调度开销).