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

13

我有一个包含16位无符号整数(MTLPixelFormatR16Uint)的MTLTexture。这些值的范围大约在7000到20000之间,其中0被用作“无数据”值,因此在下面的代码中被跳过。我想找到最小和最大值,以便将这些值重新缩放为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);
        }
    }
}

从这里我得到了一个最小值和最大值,但是对于同样的数据集,最小值和最大值通常会返回不同的值。我相当确定这是在有多个线程运行时从单个线程中获得的最小值和最大值。

第二次尝试

基于之前的尝试,这一次我正在存储每个线程的单独最小/最大值,共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组最小/最大值的数组。 我可以从中找到最小值中的最低值,但这似乎是一种不好的方法。 希望能指导正确的方向,谢谢!

2个回答

13

金属着色语言有原子比较并交换函数,您可以使用它们来比较内存位置上的现有值和一个值,并且如果它们不相等,则替换该位置上的值。使用这些函数,您可以创建一组原子比较并根据需要替换(大于|小于)的操作:

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));
}

为了使用这些,你可以创建一个缓冲区,每个线程组包含一个交错的min、max对。然后,在内核函数中从纹理中读取并有条件地写入最小值和最大值:

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)。我的理解是原子操作逐线程组应用(顺便指正一下我的任何假设,如果错了的话)? 我使用通过tpt变量中的threads_per_threadgroup注释传递到内核中的16x16线程按线程组。 我不确定这是我的线程组网格的宽度吗?例如;纹理大小为192x160,线程组网格为12x10,并且偏移计算为atomicBuffer + ((tgpig[1] * 12 + tgpig[0]) * 2) - lock
请原谅最后一行中的硬编码12。我想我试图说的是,在min_max_per_threadgroup内核中用threadgroups_per_grid替换threads_per_threadgroup是否可以解决它? - lock
@lock 是的,你完全正确。 我在实现中很幸运,因为 "threads_per_threadgroup" 恰好等于 "threadgroups_per_grid"。 已更正如上所述。 - warrenm
说实话,这可能比一开始就正确更有助于我的理解。再次感谢。 - lock
Warrenm,我很好奇,你是如何将Metal代码执行从Iris切换到Nvidia的? - Pop Flamingo

1

请试试这个。它在我这里运行良好。

kernel void grayscale_texture_minmax(texture2d<half, access::read> inTexture [[texture(0)]],
                             device atomic_uint *min_max [[buffer(0)]],
                             uint2 gid [[thread_position_in_grid]])
{
    if ((gid.x >= inTexture.get_width()) || (gid.y >= inTexture.get_height())) {
        return;
    }
    // true color to gray scaled
    const half4 inColor = inTexture.read(gid);
    const half clr = inColor.r;
    const uint intColor = uint(clamp(clr, 0.h, 1.h)*255.h);
    
    atomic_fetch_min_explicit(&min_max[0], intColor, memory_order_relaxed);
    atomic_fetch_max_explicit(&min_max[1], intColor, memory_order_relaxed);
}

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接