如何用Metal及Metal Shader Language计算16通道图像分通道均值与方差

嘿，这个问题我之前刚好处理过，咱们一步步拆解怎么用Metal Shader Language实现16通道图像的均值和方差计算～

首先得明确核心思路：计算均值和方差需要先拿到每个通道的像素值总和、像素值平方的总和，以及像素总数。因为是并行计算，我们得分两个阶段来做——先让每个线程组计算自己负责区域的局部总和，再把所有线程组的局部结果汇总成全局总和，最后推导均值和方差。

第一步：编写Metal Shader代码 #

我们需要两个Compute Kernel：一个负责计算局部区域的总和与平方和，另一个负责全局汇总。

1. 局部总和计算Kernel

这个Kernel会让每个线程处理一个像素的16个通道值，然后把局部结果存在线程组共享内存里，最后汇总到全局缓冲区：

#include <metal_stdlib>
using namespace metal;

// 第一阶段：计算每个线程组内的局部总和与平方和
kernel void computePartialSums(
    texture2d_array<float, access::read> in[[texture(0)]],
    device float* partialSums[[buffer(0)]], // 存储所有线程组的局部结果，每个线程组占32个float（16通道×2：sum+sum²）
    ushort3 gid[[thread_position_in_grid]],
    ushort tid[[thread_index_in_threadgroup]],
    ushort3 tg_size[[threads_per_threadgroup]]
) {
    const uint width = in.get_width();
    const uint height = in.get_height();
    // 跳过超出图像范围的线程
    if (gid.x >= width || gid.y >= height) {
        return;
    }

    // 线程组共享内存：存储当前组内16通道的sum和sum²
    threadgroup float tgSums[32];
    // 初始化共享内存
    if (tid < 32) {
        tgSums[tid] = 0.0;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // 读取当前像素的16个通道值（假设每个slice对应一个通道）
    float channelValues[16];
    for (int slice = 0; slice < 16; ++slice) {
        channelValues[slice] = in.read(gid.xy, slice).x;
    }

    // 将当前像素的通道值和平方值累加到共享内存
    for (int slice = 0; slice < 16; ++slice) {
        atomic_add(&tgSums[slice * 2], channelValues[slice]);
        atomic_add(&tgSums[slice * 2 + 1], channelValues[slice] * channelValues[slice]);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // 由线程0将当前线程组的结果写入全局缓冲区
    if (tid == 0) {
        uint groupId = get_group_id(0) + get_group_id(1) * get_num_groups(0);
        uint bufferOffset = groupId * 32;
        for (int i = 0; i < 32; ++i) {
            partialSums[bufferOffset + i] = tgSums[i];
        }
    }
}

2. 全局汇总与结果计算Kernel

这个Kernel会把所有线程组的局部结果累加，得到全局的总和与平方和，既可以把总和传到CPU端计算均值方差，也可以直接在Shader里完成最终计算：

// 第二阶段：汇总所有线程组的局部结果，计算全局总和
kernel void computeGlobalStats(
    device const float* partialSums[[buffer(0)]],
    device float* globalStats[[buffer(1)]], // 存储最终结果：[16个sum, 16个sum², 像素总数, 16个均值, 16个方差]
    uint tid[[thread_index_in_threadgroup]],
    uint numThreadgroups[[buffer(2)]], // 总线程组数量
    uint pixelCount[[buffer(3)]] // 图像总像素数 = width × height
) {
    // 线程组共享内存：用于归约求和
    threadgroup float tgTotal[32];
    if (tid < 32) {
        tgTotal[tid] = 0.0;
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // 每个线程处理多个局部结果元素，累加到共享内存
    for (uint i = tid; i < numThreadgroups * 32; i += tg_size.x) {
        atomic_add(&tgTotal[i % 32], partialSums[i]);
    }
    threadgroup_barrier(mem_flags::mem_threadgroup);

    // 归约求和（把共享内存里的值合并成最终总和）
    for (uint s = tg_size.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            tgTotal[tid] += tgTotal[tid + s];
        }
        threadgroup_barrier(mem_flags::mem_threadgroup);
    }

    // 线程0将全局总和写入缓冲区，同时计算均值和方差
    if (tid == 0) {
        // 写入16通道的sum和sum²
        for (int i = 0; i < 32; ++i) {
            globalStats[i] = tgTotal[i];
        }
        // 写入像素总数
        globalStats[32] = float(pixelCount);

        // 计算每个通道的均值和方差并写入缓冲区
        for (int slice = 0; slice < 16; ++slice) {
            float mean = globalStats[slice*2] / globalStats[32];
            float variance = (globalStats[slice*2+1] / globalStats[32]) - mean * mean;
            globalStats[33 + slice*2] = mean;
            globalStats[34 + slice*2] = variance;
        }
    }
}

第二步：CPU端配合流程 #

1. 准备资源：

   • 把16通道图像加载为MTLTexture，格式选MTLPixelFormatR32Float，array长度设为16（每个slice对应一个通道）。
   • 创建两个缓冲区：partialSumsBuffer（大小为线程组数量 × 32 × sizeof(float)）、globalStatsBuffer（大小至少65 × sizeof(float)，存sum、sum²、count、均值、方差）。
   • 创建两个Compute Pipeline State，分别对应上面两个Kernel。

2. 调度Kernel：

   • 先调度computePartialSums：线程组大小选16x16（可根据GPU性能调整），计算线程组数量：threadgroupsX = (width + 15)/16，threadgroupsY = (height +15)/16。
   • 再调度computeGlobalStats：线程组大小选32或64即可，因为只需要处理32个元素的归约。

3. 读取结果：

   • 从globalStatsBuffer读取数据，直接获取每个通道的均值和方差，或者自己用总和推导：

     // 示例：读取第k个通道的均值和方差
     int k = 5;
     float mean = globalStats[33 + k*2];
     float variance = globalStats[34 + k*2];
     

注意事项 #

• 线程组共享内存的大小要合理，这里用32个float完全足够（16通道×2）。
• 原子操作和线程组屏障必须正确使用，避免数据竞争。
• 如果图像尺寸极大，可以考虑分多阶段归约，但上面的两阶段方案已经能覆盖绝大多数场景。

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