前言

Metal入门教程总结

正文

核心思路

首先，我们用直方图来表示一张图像：横坐标代表的是颜色值，纵坐标代表的是该颜色值在图像中出现次数。

如图，对于某些图像，可能出现颜色值集中分布在某个区间的情况。
直方图均衡化(Histogram Equalization) ，指的是对图像的颜色值进行重新分配，使得颜色值的分布更加均匀。

本文用compute shader对图像的颜色值进行统计，然后计算得出映射关系，由fragment shader进行颜色映射处理。

效果展示

具体步骤

1、Metal的render管道、compute管道配置；

同前文，不再赘述，详见Metal入门教程总结。

2、CPU进行直方图均衡化处理；

• 2.1 把UIImage转成Bytes；

• 2.2 颜色统计；

    // CPU进行统计
    Byte *color = (Byte *)spriteData;
    for (int i = 0; i < width * height; ++i) {
        for (int j = 0; j < LY_CHANNEL_NUM; ++j) {
            uint c = color[i * 4 + j];
            ++cpuColorBuffer.channel[j][c];
        }
    }

• 2.3 映射关系；

    int rgb[3][LY_CHANNEL_SIZE], sum = (int)(width * height);
    int val[3] = {0};
    // 颜色映射
    for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < LY_CHANNEL_SIZE; ++j) {
            val[i] += cpuColorBuffer.channel[i][j];
            rgb[i][j] = val[i] * 1.0 * (LY_CHANNEL_SIZE - 1) / sum;
        }
    }

• 2.4 颜色值修改；

    // 值修改
    for (int i = 0; i < width * height; ++i) {
        for (int j = 0; j < LY_CHANNEL_NUM; ++j) {
            uint c = color[i * 4 + j];
            color[i * 4 + j] = rgb[j][c];
        }
    }

最后用处理之后的Bytes生成新图片。

3 GPU进行直方图均衡化处理；

• 3.1 compute shader进行颜色统计；

kernel void
grayKernel(texture2d<float, access::read>  sourceTexture  [[textureLYKernelTextureIndexSource]], // 纹理输入，
           device LYColorBuffer &out [[buffer(LYKernelBufferIndexOutput)]], // 输出的buffer
           uint2                          grid         [[thread_position_in_grid]]) // 格子索引
{
    // 边界保护
    if(grid.x < sourceTexture.get_width() && grid.y < sourceTexture.get_height())
    {
        float4 color  = sourceTexture.read(grid); // 初始颜色
        int3 rgb = int3(color.rgb * SIZE); // 乘以SIZE，得到[0, 255]的颜色值
        // 颜色统计，每个像素点计一次
        atomic_fetch_add_explicit(&out.channel[0][rgb.r], 1, memory_order_relaxed);
        atomic_fetch_add_explicit(&out.channel[1][rgb.g], 1, memory_order_relaxed);
        atomic_fetch_add_explicit(&out.channel[2][rgb.b], 1, memory_order_relaxed);
    }
}

atomic_fetch_add_explicit是用于在多线程进行数据操作，具体的函数解释见这里。

• 3.2 映射关系处理；
  compute shader回调后，根据GPU统计的颜色分布结果，求出映射关系；

        LYLocalBuffer *buffer = (LYLocalBuffer *)strongSelf.colorBuffer.contents; // GPU统计的结果
        LYLocalBuffer *convertBuffer = self.convertBuffer.contents; // 颜色转换的buffer
        int sum = (int)(self.sourceTexture.width * self.sourceTexture.height); // 总的像素点
        int val[3] = {0}; // 累计和
        for (int i = 0; i < 3; ++i) {
            for (int j = 0; j < LY_CHANNEL_SIZE; ++j) {
                val[i] += buffer->channel[i][j]; // 当前[0, j]累计出现的总次数
                convertBuffer->channel[i][j] = val[i] * 1.0 * (LY_CHANNEL_SIZE - 1) / sum;
                
                // 对比CPU和GPU处理的结果
                if (buffer->channel[i][j] != strongSelf->cpuColorBuffer.channel[i][j]) {
                    // 如果不相同，则把对应的结果输出
                    printf("%d, %d, gpuBuffer:%u  cpuBuffer:%u \n", i, j, buffer->channel[i][j], strongSelf->cpuColorBuffer.channel[i][j]);
                }
            }
        }
        memset(buffer, 0, strongSelf.colorBuffer.length);

3.3 根据映射关系处理原图片，并渲染到屏幕上；

fragment float4
samplingShader(RasterizerData input [[stage_in]], // stage_in表示这个数据来自光栅化。（光栅化是顶点处理之后的步骤，业务层无法修改）
               texture2d<float> colorTexture [[ texture(LYFragmentTextureIndexSource) ]], // texture表明是纹理数据，LYFragmentTextureIndexSource是索引
               device LYLocalBuffer &convertBuffer [[buffer(LYFragmentBufferIndexConvert)]]) // 转换的buffer
{
    constexpr sampler textureSampler (mag_filter::linear, min_filter::linear); // sampler是采样器
    float4 colorSample = colorTexture.sample(textureSampler, input.textureCoordinate); // 得到纹理对应位置的颜色
    int3 rgb = int3(colorSample.rgb * SIZE); // 记得先乘以SIZE
    colorSample.rgb = float3(convertBuffer.channel[0][rgb.r], convertBuffer.channel[1][rgb.g], convertBuffer.channel[2][rgb.b]) / SIZE; // 返回的值也要经过归一化处理
    return colorSample;
}

遇到的问题

1、统计结果集中在头部

问题表现：
统计结果异常，集中在前面两个值。
如下，green通道的颜色集中在r[0]和r[1]上：

// 0~255颜色值的分布
28269 4492 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0

问题分析：
实际上，gpu里面存着的是0.0~1.0的值；（归一化）
统计的值全部是在前面，是因为没有乘以255！
先用CPU实现了直方图均衡化，在实现shader的时候，参考CPU的代码实现，犯了这个错误。

2、cpu和gpu统计结果相差较多

问题表现：
如下代码，buffer是gpu统计的颜色值分布结果，cpuColorBuffer是cpu统计的颜色值分布结果。
理论上结果应该接近，但实际上printf出来的差异非常多。

        for (int i = 0; i < 3; ++i) {
            for (int j = 0; j < LY_CHANNEL_SIZE; ++j) {
                val[i] += buffer->channel[i][j];
                convertBuffer->channel[i][j] = val[i] * 1.0 * (LY_CHANNEL_SIZE - 1) / sum;
                
                // 对比CPU和GPU处理的结果
                if (buffer->channel[i][j] != strongSelf->cpuColorBuffer.channel[i][j]) {
                    // 如果不相同，则把对应的结果输出
                    printf("%d, %d, gpuBuffer:%u  cpuBuffer:%u \n", i, j, buffer->channel[i][j], strongSelf->cpuColorBuffer.channel[i][j]);
                }
            }
        }

问题分析：
通过检查代码，先判定cpu统计的结果是正常。（cpu的处理过程就是正常的for循环，不易出错）
仔细观察log的不同：
0, 1, gpuBuffer:763 cpuBuffer:762
结果很接近，但是有细微的差距。
我们知道gpu是浮点数的处理，而cpu是整数型处理，浮点数到整数中间有精度的问题。
此时再看我们的shader，我们是以half来进行计算，这样统计出来的结果会有点误差。

grayKernel(texture2d<half, access::read>  sourceTexture  [[texture(LYFragmentTextureIndexTextureSource)]],
           device LYColorBuffer &out [[buffer(LYKernelBufferIndexOutput)]],
           uint2                          grid         [[thread_position_in_grid]])

通过把精度从half改成float，cpu和gpu的统计差异就只有3个：

0, 248, gpuBuffer:23215  cpuBuffer:22854
1, 74, gpuBuffer:23201  cpuBuffer:22840
2, 64, gpuBuffer:23336  cpuBuffer:22975

3、gpu渲染的图片为白色

问题表现：
在gpu统计的结果与cpu接近的情况下，把映射buffer传给fragment shader，最后进行一次颜色处理。
但是结果是白色的图片，shader的代码如下：

fragment float4
samplingShader(RasterizerData input [[stage_in]], // stage_in表示这个数据来自光栅化。（光栅化是顶点处理之后的步骤，业务层无法修改）
               texture2d<float> colorTexture [[ texture(LYFragmentTextureIndexTextureSource) ]], // texture表明是纹理数据，LYFragmentTextureIndexTextureSource是索引
               device LYLocalBuffer &localBuffer [[buffer(LYFragmentBufferIndexConvert)]])
{
    constexpr sampler textureSampler (mag_filter::linear,
                                      min_filter::linear); // sampler是采样器
    
    float4 colorSample = colorTexture.sample(textureSampler, input.textureCoordinate); // 得到纹理对应位置的颜色
    int3 rgb = int3(colorSample.rgb);
    colorSample.rgb = float3(localBuffer.channel[0][rgb.r], localBuffer.channel[1][rgb.g], localBuffer.channel[2][rgb.b]);
    
    return colorSample;
}

问题分析：
我们先把colorSample.rgb = ...的这行代码屏蔽，发现渲染结果是正常的，那么问题就出现在映射处理上面。
再通过Xcode的Capture GPU Frame工具，查看传入的映射buffer数据，也是正常的数据。
那么问题可能出现int3 rgb的初始化，或者从映射buffer读取数据。
观察到int3 rgb = int3(colorSample.rgb)，是有一个float->int的操作，联想到前面提到的归一化处理，马上明白：在这里的初始化时应该乘以SIZE。

那么问题是否就此解决？不是的。

我们在进行颜色转换的时候，float->int 需要乘以SIZE；
在获取到映射buffer里面对应颜色的值后，仍需要做一次int->float的处理，除以SIZE；
如果下：

    float4 colorSample = colorTexture.sample(textureSampler, input.textureCoordinate); // 得到纹理对应位置的颜色
    int3 rgb = int3(colorSample.rgb * size);
    colorSample.rgb = float3(localBuffer.channel[0][rgb.r], localBuffer.channel[1][rgb.g], localBuffer.channel[2][rgb.b]) / size;
    

4、映射结果异常

问题表现：
问题如下，映射结果应该是0~255的值，但是通过Xcode看到最终的映射值远超过255，甚至接近255*2的数字。

问题分析：
下面是映射的算法

    int rgb[3][LY_CHANNEL_SIZE], sum = (int)(width * height);
    int val[3] = {0};
    for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < LY_CHANNEL_SIZE; ++j) {
            val[i] += cpuColorBuffer.channel[i][j];
            rgb[i][j] = val[i] * 1.0 * (LY_CHANNEL_SIZE - 1) / sum;
        }
    }

sum是固定值，LY_CHANNEL_SIZE是常量值256，那么映射结果超过255的原因就是val[i]的统计结果太大！
通过Xcode调试，确实如此：

那么，会是什么原因导致？

在看到结果接近255的两倍时，大概猜测可能是重复运算导致。
我们的均衡化处理是在MTKView的回调进行，如下：

- (void)drawInMTKView:(MTKView *)view {
    [self customDraw];
}

这里会回调多次，从而导致多次执行compute shader的颜色统计，这里可以引入isDrawing的临时变量解决：

- (void)drawInMTKView:(MTKView *)view {
    if (!self.isDrawing) {
        self.isDrawing = YES;
        [self customDraw];
    }
}

但是，问题并没有彻底解决：首次统计正常，但是第二次处理的时候就会累积上一次的统计值。
如何对值进行清理？（这里实际上只处理一次也行，但是debug过程中需要通过Xcode的GPU Capture Frame工具进行查看，而这个工具需要多次渲染）
我们知道MTLBuffer是cpu、gpu都可以操作的buffer，那么在cpu直接清除这个数据即可。
在 commandBuffer addCompletedHandler:^(){}的结束回调中，使用memset(buffer, 0, strongSelf.colorBuffer.length)清理统计结果。

5、映射结果最大值为256

问题表现：
在踩过上面的各种坑之后，直方图均衡化的效果也已经展现，但是仍有一点小问题：
映射结果buffer的数字范围是0~256，而不是255。

问题分析：
根据直方图均衡化的算法，我们知道是因为像素颜色值的统计，结果稍微偏大。
回顾Compute Shader的代码：

kernel void
grayKernel(texture2d<float, access::read>  sourceTexture  [[texture(LYFragmentTextureIndexTextureSource)]],
           device LYColorBuffer &out [[buffer(LYKernelBufferIndexOutput)]],
           uint2                          grid         [[thread_position_in_grid]]) {
    // 边界保护
    if(grid.x <= sourceTexture.get_width() && grid.y <= sourceTexture.get_height())
    {
        float4 color  = sourceTexture.read(grid); // 初始颜色
        int3 rgb = int3(color.rgb * size); //
        atomic_fetch_add_explicit(&out.channel[0][rgb.r], 1, memory_order_relaxed);
        atomic_fetch_add_explicit(&out.channel[1][rgb.g], 1, memory_order_relaxed);
        atomic_fetch_add_explicit(&out.channel[2][rgb.b], 1, memory_order_relaxed);
    }
}

Metal的Compute Shader是按每组网格进行处理，那么可能会出现边界超过图像的情况，所以添加了边界保护。
但是，这里存在误判的情况：边界判断不应该是<=，而是<。
因为这个原因，会导致统计的结果偏大，最终出现256的情况。

在修复完这个问题后，Demo再无其他问题，GPU的处理结果也和CPU的处理结果完全一致！

总结

本文是在前文的Metal入门教程基础上进行更复杂的尝试，过程中也遇到较多问题，最终demo也顺利完成，地址在Github。

——-2018年8月31日 周五 23:51，踩着8月的尾巴更新这一篇文章，九月是休假月。