Metal 进阶：计算通道

一牛
336 阅读3分钟

引言

大家好，我是一牛，很高兴和大家又见面了。我们在上一篇文章中介绍了离屏渲染。通过创建两个渲染通道，一个负责离屏渲染，另一个负责将离屏渲染的结果渲染到窗口。今天我分享的是Metal 中的计算通道。在图形编程中，计算通道（Compute Pass） 指的是GPU 计算管线中的一次执行过程，它专门用于计算任务（如物理模拟、图像处理、AI 推理等），而不是传统的光栅化渲染。在本例中，我将使用计算通道实现图片处理 - 灰度化图片并渲染到窗口。

计算通道

计算函数

// Grayscale compute kernel
constant half3 kRec709Luma = half3(0.2126, 0.7152, 0.0722);
kernel void
grayscaleKernel(texture2d<half, access::read>  inTexture  [[texture(0)]],
                texture2d<half, access::write> outTexture [[texture(1)]],
                uint2                          gid        [[thread_position_in_grid]])
{   
    half4 inColor  = inTexture.read(gid);
    half  gray     = dot(inColor.rgb, kRec709Luma);
    outTexture.write(half4(gray, gray, gray, 1.0), gid);
}

和顶点着色器、片段着色器类似，计算函数是由关键字kernel 修饰的，所以它也被称为内核函数。这个函数是将输入的inTexture转换成灰度图片，将输出写入到outTexture。这里需要注意的是，inputTexture是只读，而outputTexture是可写的。thread_position_in_grid 在本例只指代的就是图片的像素坐标。

计算管线

在Metal 中创建计算管线非常容易,绑定计算函数即可。

guard let kernelFunction = library?.makeFunction(name: "grayscaleKernel") else {
    fatalError("fail to load kernel function - grayscaleKernel")
}
do {
    computePipelineState = try device.makeComputePipelineState(function: kernelFunction)
} catch {
    fatalError("fail to create compute pipelineState")
}

编码计算命令

let computeEncoder = commanderBuffer?.makeComputeCommandEncoder()
computeEncoder?.setComputePipelineState(computePipelineState)
computeEncoder?.setTexture(inputTexture, index: 0)
computeEncoder?.setTexture(outputTexture, index: 1)
let threadsPerGrid = MTLSize(width: inputTexture.width, height: inputTexture.height, depth: 1)
computeEncoder?.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)
computeEncoder?.endEncoding()

var threadsPerThreadgroup: MTLSize {
    let w = computePipelineState?.threadExecutionWidth ?? 0
    let h = (computePipelineState?.maxTotalThreadsPerThreadgroup ?? 0) / w
    return MTLSize(width: w, height: h, depth: 1)
}

这里threadsPerGrid和图片的尺寸一致，threadsPerThreadgroup是一个线程组。如图，图片被分割成了多个线程组，而线程组大小并不总是一致。

media-2922064~dark@2x.png

这里还需要注意的是，由于outputTexture对应计算函数的outTexture, 它需要可写。在本例中，我们还需要将它渲染到窗口，所以它还需要可读，所以我们在创建纹理的时候需要申明成可读可写。

func createOutputTextureDescriptor(inputTexture: MTLTexture) -> MTLTextureDescriptor {
    let descriptor = MTLTextureDescriptor()
    descriptor.textureType = .type2D
    descriptor.width = inputTexture.width
    descriptor.height = inputTexture.height
    descriptor.pixelFormat = inputTexture.pixelFormat
    // shaderWrite -> compute pass, shaderRead -> render pass
    descriptor.usage = [.shaderWrite, .shaderRead]
    return descriptor
}

总结

通过 计算通道，我们可以在 Metal 中高效地处理各种计算任务，如图像滤波、深度学习推理、物理模拟等。理解和掌握计算通道的使用，不仅能够帮助我们优化图形渲染流程，还能拓展应用程序的计算能力，尤其在图像处理和 AI 加速方面具有广泛应用。

希望大家能通过本文了解 Metal 中的计算通道，并掌握其在图像处理中的应用。如果你有任何问题或想进一步讨论，欢迎留言或分享你的看法！

本例源码

avatar
文章
阅读
粉丝