如何在GPU上通过掩码值/纹理计算平均值？

Question

我想做的是：

RWTexture<float3>TexColor; // RBGA32
Texture<uint> TexMasked;   // R32

void get_masked_mean(uint2 pixelPos)
{
    //? How do it on GPU?
    // TexColor[pixelPos] = calculateAverage(TexColor[pixel]) for all
    // pixels satisfying (TexMasked[pixel] == TexMasked[pixel])
}

我尝试将 TexColor 和 TexMasked 读取到 CPU，但它非常解决：

std::unordered_map<uint, float4> resultMap;
std::vector<std::vector<float3>> colorData[][]  = TexColor.readbakToCPU(); // very slow
std::vector<std::vector<uint>>   maskedData[][] = TexMasked.readbakcToCPU();

for (int r = 0; r < rowCount; ++r)     // 1920
{
    for (int c = 0; c < colCount; ++c) // 1080
    {
       uint key = maskedData[r][c];
       uint oldCount = resultMap[key].w;
       uint newCount = oldCount + 1;

       // very slow
       resultMap[key] = float4((colorData[r][c] + count * resultMap[key].xyz) / newCount), newCount);
    }
}

for (int r = 0; r < rowCount; ++r)
{
    for (int c = 0; c < colCount; ++c)
    {
       
       uint key = maskedData[r][c];
       colorData[r][c] = resultMap[key].xyz; // very solw
    }
}

TexColor.uploadToGPU(colorData);

从 GPU 读取数据到 CPU，再加上约 1,000,000 次哈希查询非常非常慢。这里有更高效的CPU代码吗？或者基于 GPU 的代码（计算着色器）？

Answer 1

嗯，有很多方法可以优化这段代码。

如果你想坚持使用基于 CPU 的计算（这更容易做到，但也总是会明显变慢或伴随其他吸引力），我要做的第一件事就是找到更快的哈希映射实现。您正在执行大约 4,000,000 次哈希表查找，因此即使哈希表访问速度的微小改进也可以产生可衡量的差异。这不会是一个巨大的收获，但它是一个相当容易实现的目标。

如果您多次执行此计算（例如，在交互式应用程序中渲染单个帧时），您还可以将工作流水线化：您始终访问几帧前的内存，而不是立即访问纹理内存，并且GPU 随后还会访问几帧前上传的内存。这会大大增加内存成本（因为您必须始终保持多个帧处于飞行状态），并且 GPU 始终只能看到几帧前的计算结果，但这也将完全消除读取的性能成本如果您保持足够数量的飞行帧，则 CPU 上的纹理并将结果上传回 GPU。显然，如果这是一次性计算，或者如果 GPU 上的延迟结果不可接受，则这不是一种选择。

然后，您可以通过使用多个线程进一步优化代码，尽管这项工作可能不值得。

在 GPU 上进行此计算是可能的，但复杂得多。由于 GPU 是高度并行化的，因此在 GPU 上实现哈希映射非常具有挑战性。相反，我会使用多个着色器来重新组织数据，直到您不再需要哈希表为止。基本思想是，按键对数据进行排序（存储排序数据中每个条目的原始位置），然后使用前缀和将具有相同键的所有连续条目“展平”为单个数字，然后您可以用于为每个唯一键安排工作。它的确切实现取决于一些细节 - 例如，您可以使用间接执行来为每个唯一键安排一个线程（或一个线程组/等），或者您可以重写您的数学，以便它可以原子地完成，然后每个像素运行一个线程（请注意，仅某些供应商特定的扩展支持浮点上的原子操作，并且并非在每个硬件上都可用）。 最后，使用排序前存储的每个条目的原始位置，以及该条目在排序后的数据中的位置，您可以将结果写入每个像素。

整个过程需要多个计算着色器和额外的缓冲区，并且以快速高效的方式编写计算着色器和整个管道是它自己的一个主题。但如果你投入工作和时间，你可以比任何 CPU 端方法更快（我猜在好的硬件上花费的时间可能不到 1 毫秒）。