如何在Swift中高效生成指定尺寸的高斯分布MTLTexture/MTLBuffer？

问题背景 #

开发实时视频滤镜应用时，需要生成单变量高斯分布的随机Buffer/Texture。Python中用numpy实现的代码耗时约0.15秒：

h = 1170
w = 2532

with Timer():
    noise = np.random.normal(size=w * h * 3)

plt.imshow(noise.reshape(w,h,3))
plt.show()

但自行编写的Swift代码性能极差，无法满足实时需求，核心问题在于循环内重复初始化随机源和分布，导致大量不必要的开销。

原Swift代码的性能瓶颈 #

• 嵌套循环中每次创建GKRandomSource和GKGaussianDistribution，初始化开销极大
• CPU串行遍历每个像素，无并行优化
• 循环内重复计算常量值，浪费算力

优化方案 #

1. 快速修复：CPU串行版本优化 #

核心是复用随机源和高斯分布，提前计算常量：

private func generateNoiseTextureBuffer(width: Int, height: Int) -> [Float] {
    let w = Float(width)
    let h = Float(height)
    let pixelCount = width * height
    var noiseData = [Float](repeating: 0, count: pixelCount * 4)
    
    // 仅初始化一次随机源与高斯分布
    let randomSource = GKRandomSource()
    let gaussianDist = GKGaussianDistribution(randomSource: randomSource, mean: 0.0, deviation: 1.0)
    
    // 预计算所有常量，避免循环内重复计算
    let scale = sqrt(2.0 * min(w, h) * (2.0 / Float.pi))
    let maxX = w - 1.0
    let maxY = h - 1.0
    
    for yi in 0..<height {
        for xi in 0..<width {
            let pixelIndex = yi * width + xi
            let x = Float(xi)
            let y = Float(yi)
            
            // 复用同一个分布生成随机数
            let randX = Float(gaussianDist.nextUniform())
            let randY = Float(gaussianDist.nextUniform())
            
            let rx = floor(max(min(x + scale * randX, maxX), 0.0))
            let ry = floor(max(min(y + scale * randY, maxY), 0.0))
            
            let dataIndex = pixelIndex * 4
            noiseData[dataIndex] = rx + 0.5
            noiseData[dataIndex + 1] = ry + 0.5
            noiseData[dataIndex + 2] = 1.0
            noiseData[dataIndex + 3] = 1.0
        }
    }
    return noiseData
}

2. 进一步提速：SIMD并行优化 #

利用Swift的SIMD类型一次性生成多个随机数，减少循环次数：

// 批量生成4个高斯随机数
private func generateSIMDGaussianNumbers(_ dist: GKGaussianDistribution) -> SIMD4<Float> {
    return SIMD4(
        Float(dist.nextUniform()),
        Float(dist.nextUniform()),
        Float(dist.nextUniform()),
        Float(dist.nextUniform())
    )
}

private func generateNoiseTextureBufferSIMD(width: Int, height: Int) -> [Float] {
    let w = Float(width)
    let h = Float(height)
    let pixelCount = width * height
    var noiseData = [Float](repeating: 0, count: pixelCount * 4)
    
    let randomSource = GKRandomSource()
    let gaussianDist = GKGaussianDistribution(randomSource: randomSource, mean: 0.0, deviation: 1.0)
    let scale = sqrt(2.0 * min(w, h) * (2.0 / Float.pi))
    let maxX = w - 1.0
    let maxY = h - 1.0
    
    let batchSize = 4
    let remainingPixels = pixelCount % batchSize
    
    // 批量处理像素
    for i in stride(from: 0, to: pixelCount - remainingPixels, by: batchSize) {
        let simdRandsX = generateSIMDGaussianNumbers(gaussianDist)
        let simdRandsY = generateSIMDGaussianNumbers(gaussianDist)
        
        for j in 0..<batchSize {
            let pixelIndex = i + j
            let xi = pixelIndex % width
            let yi = pixelIndex / width
            let x = Float(xi)
            let y = Float(yi)
            
            let rx = floor(max(min(x + scale * simdRandsX[j], maxX), 0.0))
            let ry = floor(max(min(y + scale * simdRandsY[j], maxY), 0.0))
            
            let dataIndex = pixelIndex * 4
            noiseData[dataIndex] = rx + 0.5
            noiseData[dataIndex + 1] = ry + 0.5
            noiseData[dataIndex + 2] = 1.0
            noiseData[dataIndex + 3] = 1.0
        }
    }
    
    // 处理剩余不足batchSize的像素
    for i in (pixelCount - remainingPixels)..<pixelCount {
        let xi = i % width
        let yi = i / width
        let x = Float(xi)
        let y = Float(yi)
        
        let randX = Float(gaussianDist.nextUniform())
        let randY = Float(gaussianDist.nextUniform())
        
        let rx = floor(max(min(x + scale * randX, maxX), 0.0))
        let ry = floor(max(min(y + scale * randY, maxY), 0.0))
        
        let dataIndex = i * 4
        noiseData[dataIndex] = rx + 0.5
        noiseData[dataIndex + 1] = ry + 0.5
        noiseData[dataIndex + 2] = 1.0
        noiseData[dataIndex + 3] = 1.0
    }
    
    return noiseData
}

3. 实时场景最优解：Metal GPU生成 #

视频滤镜本身基于Metal，直接在GPU上并行生成高斯噪声，完全满足实时要求：

Metal Shader代码（存为.metal文件）

#include <metal_stdlib>
using namespace metal;

// Box-Muller变换生成高斯分布随机数
float gaussianRandom(uint2 seed) {
    float u1 = fract(sin(dot(seed, uint2(1299792, 344994))) * 43758.5453);
    float u2 = fract(sin(dot(seed, uint2(987654, 654321))) * 789012.3456);
    return sqrt(-2.0 * log(u1)) * cos(2.0 * M_PI_F * u2);
}

kernel void generateGaussianNoise(texture2d<float, access::write> outputTexture [[texture(0)]],
                                  uint2 gid [[thread_position_in_grid]]) {
    int width = outputTexture.get_width();
    int height = outputTexture.get_height();
    float w = float(width);
    float h = float(height);
    
    float scale = sqrt(2.0 * min(w, h) * (2.0 / M_PI_F));
    float maxX = w - 1.0;
    float maxY = h - 1.0;
    
    float x = float(gid.x);
    float y = float(gid.y);
    
    // 用不同种子生成独立的随机数
    float randX = gaussianRandom(gid);
    float randY = gaussianRandom(gid ^ uint2(123, 456));
    
    float rx = floor(clamp(x + scale * randX, 0.0, maxX));
    float ry = floor(clamp(y + scale * randY, 0.0, maxY));
    
    outputTexture.write(float4(rx + 0.5, ry + 0.5, 1.0, 1.0), gid);
}

Swift端调用代码

import MetalKit

class GaussianNoiseGenerator {
    private let device: MTLDevice
    private let commandQueue: MTLCommandQueue
    private let computePipelineState: MTLComputePipelineState
    
    init?(device: MTLDevice) {
        self.device = device
        self.commandQueue = device.makeCommandQueue()!
        
        // 加载默认Metal库
        guard let library = device.makeDefaultLibrary() else { return nil }
        guard let kernelFunction = library.makeFunction(name: "generateGaussianNoise") else { return nil }
        
        do {
            self.computePipelineState = try device.makeComputePipelineState(function: kernelFunction)
        } catch {
            print("创建Pipeline State失败: \(error)")
            return nil
        }
    }
    
    func generateNoiseTexture(width: Int, height: Int) -> MTLTexture? {
        // 创建输出纹理描述
        let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: .rgba32Float,
            width: width,
            height: height,
            mipmapped: false
        )
        textureDescriptor.usage = [.shaderWrite, .shaderRead]
        guard let outputTexture = device.makeTexture(descriptor: textureDescriptor) else { return nil }
        
        // 创建命令缓冲区与编码器
        guard let commandBuffer = commandQueue.makeCommandBuffer() else { return nil }
        guard let computeEncoder = commandBuffer.makeComputeCommandEncoder() else { return nil }
        
        computeEncoder.setComputePipelineState(computePipelineState)
        computeEncoder.setTexture(outputTexture, index: 0)
        
        // 设置线程组（16x16是Metal常用的线程组大小）
        let threadGroupSize = MTLSize(width: 16, height: 16, depth: 1)
        let threadGroups = MTLSize(
            width: (width + threadGroupSize.width - 1) / threadGroupSize.width,
            height: (height + threadGroupSize.height - 1) / threadGroupSize.height,
            depth: 1
        )
        computeEncoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadGroupSize)
        
        computeEncoder.endEncoding()
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()
        
        return outputTexture
    }
}

// 使用示例
if let device = MTLCreateSystemDefaultDevice(),
   let noiseGenerator = GaussianNoiseGenerator(device: device) {
    let noiseTexture = noiseGenerator.generateNoiseTexture(width: 2532, height: 1170)
    // 直接将noiseTexture用于视频滤镜的后续渲染流程
}

方案选择建议 #

• 实时视频滤镜场景：优先选择Metal GPU方案，并行计算性能远超CPU，且与现有Metal渲染流程无缝集成
• 非实时或快速验证：先使用优化后的CPU串行版本，SIMD版本可作为进一步提速的选项

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