基于GPU的蒙版纹理均值计算方案问询（RTX实时光照烘焙）

我在Windows平台用NVIDIA RTX显卡开发一款基于GPU的实时预览光照烘焙器，处理待烘焙场景时会做以下操作：

• 将屏幕空间光照数据写入1920x1080的RWTexture2D<float3> TexColor，作为几何表面无限精度光照贴图分辨率的结果；
• 将由objectId和uint2(lightmapUV.x, lightmapUV.y)组合而成、用于唯一标识光照贴图纹素的MaskedID写入1920x1080的RWTexture2D<uint> TexMasked，用来过滤TexColor以计算指定光照贴图分辨率的结果。

最终要把结果应用到场景提供光照贴图预览，我需要实现的核心逻辑如下（GPU端伪代码）：

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

void get_masked_mean(uint2 pixelPos)
{
    // 如何在GPU上实现？
    // TexColor[pixelPos] = 所有满足TexMasked[pixel] == TexMasked[pixelPos]的TexColor[pixel]的平均值
}

我尝试过把数据读到CPU计算，但速度极慢：

std::unordered_map<uint, float4> resultMap;
std::vector<std::vector<float3>> colorData = TexColor.readbackToCPU(); // 极慢
std::vector<std::vector<uint>> maskedData = TexMasked.readbackToCPU();

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;

       // 极慢
       resultMap[key] = float4((colorData[r][c] + oldCount * 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; // 极慢
    }
}

TexColor.uploadToGPU(colorData);

GPU到CPU的数据传输加上约100万次哈希查询导致速度极差，请问有没有更高效的CPU实现，或者基于GPU的计算着色器方案？

一、GPU计算着色器方案（推荐） #

直接在GPU上完成分组求平均，避免数据往返，是最适合实时预览的方案。可以用并行归约+分组统计的思路，分两步实现：

步骤1：分组累加统计 #

首先创建两个临时RWTexture：

• RWTexture2D<float4> TexSumCount：存储每个MaskedID对应的颜色总和（xyz）和像素数量（w）
• （可选）RWTexture2D<uint> TexVisited：标记该MaskedID是否已被处理，避免重复累加

计算着色器代码（HLSL）：

Texture2D<float3> TexColor;
Texture2D<uint> TexMasked;
RWTexture2D<float4> TexSumCount; // 初始化全0
RWTexture2D<uint> TexVisited;    // 初始化全0

[numthreads(16,16,1)]
void CS_Accumulate(uint3 dispatchThreadID : SV_DispatchThreadID)
{
    uint2 pixelPos = dispatchThreadID.xy;
    if (pixelPos.x >= 1920 || pixelPos.y >= 1080) return;

    uint maskID = TexMasked[pixelPos];
    float3 color = TexColor[pixelPos];

    // 原子操作确保多线程累加安全
    InterlockedAdd(TexSumCount[maskID].x, color.x);
    InterlockedAdd(TexSumCount[maskID].y, color.y);
    InterlockedAdd(TexSumCount[maskID].z, color.z);
    InterlockedAdd(TexSumCount[maskID].w, 1);
}

注意：如果MaskedID的取值范围很大（超过纹理尺寸），不要用纹理存储SumCount，改用RWStructuredBuffer<float4>，并配合原子计数器分配唯一索引——比如先遍历所有像素收集唯一MaskedID，再分配索引，或者直接用原子递增的方式动态分配。

步骤2：计算并写入平均值 #

再启动一个计算着色器，遍历每个像素，根据MaskedID从SumCount中取出平均值，写入TexColor：

RWTexture2D<float3> TexColor;
Texture2D<uint> TexMasked;
Texture2D<float4> TexSumCount;

[numthreads(16,16,1)]
void CS_WriteAverage(uint3 dispatchThreadID : SV_DispatchThreadID)
{
    uint2 pixelPos = dispatchThreadID.xy;
    if (pixelPos.x >= 1920 || pixelPos.y >= 1080) return;

    uint maskID = TexMasked[pixelPos];
    float4 sumCount = TexSumCount[maskID];
    if (sumCount.w == 0) return; // 避免除以0

    TexColor[pixelPos] = sumCount.xyz / sumCount.w;
}

优化点 #

• 使用numthreads(16,16,1)或32,32,1这类适配RTX显卡的线程组大小，提升缓存命中率；
• 如果MaskedID是连续的，直接用纹理存储SumCount效率最高；如果不连续，用结构化缓冲区配合原子操作；
• 可将两步合并为一个计算着色器，但分开实现更清晰且便于调试。

二、优化后的CPU方案 #

如果必须用CPU计算，核心优化方向是减少数据传输开销和哈希表性能损耗：

1. 减少数据传输 #

• 直接映射GPU内存到CPU，避免多轮拷贝：用Map方法获取纹理的连续内存指针，读取/修改完成后Unmap自动同步回GPU，不要拆分成二维vector；
• 示例：float3* colorData = (float3*)TexColor.Map(D3D11_MAP_READ_WRITE);，uint* maskedData = (uint*)TexMasked.Map(D3D11_MAP_READ);

2. 优化哈希表性能 #

• 提前预留足够空间：resultMap.reserve(estimatedUniqueIDs);，避免哈希表频繁扩容；
• 用std::flat_hash_map（Abseil/Boost库）代替std::unordered_map，前者缓存友好性更强，查询和插入速度更快；
• 若MaskedID取值范围已知且连续，直接用std::vector<float4>代替哈希表，索引即为MaskedID，完全规避哈希查询开销。

优化后的CPU代码示例 #

// 直接映射GPU内存，避免拷贝
float3* colorData = nullptr;
uint* maskedData = nullptr;
TexColor.Map(D3D11_MAP_READ_WRITE, (void**)&colorData);
TexMasked.Map(D3D11_MAP_READ, (void**)&maskedData);

int totalPixels = 1920 * 1080;
std::flat_hash_map<uint, float4> resultMap;
resultMap.reserve(10000); // 预估计唯一ID数量

for (int i = 0; i < totalPixels; ++i)
{
    uint key = maskedData[i];
    auto it = resultMap.find(key);
    if (it != resultMap.end())
    {
        float4& sumCount = it->second;
        sumCount.xyz += colorData[i];
        sumCount.w += 1;
    }
    else
    {
        resultMap[key] = float4(colorData[i], 1.0f);
    }
}

// 计算平均值并写回
for (int i = 0; i < totalPixels; ++i)
{
    uint key = maskedData[i];
    float4 sumCount = resultMap[key];
    colorData[i] = sumCount.xyz / sumCount.w;
}

// 解除映射，自动同步回GPU
TexColor.Unmap();
TexMasked.Unmap();

CPU方案局限性 #

即使优化后，CPU处理速度仍远不及GPU，尤其不适合实时预览场景，建议优先采用GPU方案。

内容的提问来源于stack exchange，提问作者carnot