本文件通過Metal compute shader對攝像頭當前捕獲的畫面進行簡單的Gamma校正，繪製到螢幕（MTKView）及將渲染結果儲存成UIImage。文件最後簡要討論了Metal compute shader的dispatchThreadgroups配置問題。

文件結構：

1. 配置AVCaptureSession獲取攝像頭當前畫面
2. 初始化Compute Shader環境
3. 編寫Gamma校正shader程式碼
4. 渲染Compute Shader處理後的紋理到螢幕
5. 讀取Metal渲染結果並生成UIImage
6. 討論：Metal compute shader合理的dispatchThreadgroups設定

1. 配置AVCaptureSession獲取攝像頭當前畫面

參考我之前的文件iOS VideoToolbox硬編H.265（HEVC）H.264（AVC）：1 概述進行攝像頭的配置，簡單起見，令攝像頭輸出畫面為豎直方向的RGBA資料，後續文件再實踐Metal Shader實現YUV轉RGB，然後進行各種濾鏡的疊加，參考程式碼如下。

let device = AVCaptureDevice.defaultDevice(withMediaType: AVMediaTypeVideo)
let input = try? AVCaptureDeviceInput(device: device)
if session.canAddInput(input) {
    session.addInput(input)
}

let output = AVCaptureVideoDataOutput()
output.videoSettings = [kCVPixelBufferPixelFormatTypeKey as AnyHashable : kCVPixelFormatType_32BGRA]
output.setSampleBufferDelegate(self, queue: DispatchQueue(label: "CamOutputQueue"))
if session.canAddOutput(output) {
    session.addOutput(output)
}

if session.canSetSessionPreset(AVCaptureSessionPreset1920x1080) {
    session.canSetSessionPreset(AVCaptureSessionPreset1920x1080)
}

session.beginConfiguration()

for (_, connection) in output.connections.enumerated() {
    for (_, port) in (connection as! AVCaptureConnection).inputPorts.enumerated() {
        if (port as! AVCaptureInputPort).mediaType == AVMediaTypeVideo {
            videoConnection = connection as? AVCaptureConnection
            break
        }
    }
    if videoConnection != nil {
        break;
    }
}

if (videoConnection?.isVideoOrientationSupported)! {
    videoConnection?.videoOrientation = .portrait
}

session.commitConfiguration()
session.startRunning()

2. 初始化Compute Shader環境

Core Video給Metal提供了類似OpenGL ES建立紋理的介面CVMetalTextureCache。除此之外，還需進行Metal要求的MTLLibrary等準備工作，參考程式碼如下。

var textureCache : CVMetalTextureCache?
var imageTexture: MTLTexture?

var commandQueue: MTLCommandQueue?
var library: MTLLibrary?
var pipeline: MTLComputePipelineState?

//------------
device = MTLCreateSystemDefaultDevice()

mtlView.device = device
mtlView.framebufferOnly = false
mtlView.clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 1)

library = device?.newDefaultLibrary()
guard let function = library?.makeFunction(name: "gamma_filter") else {
    fatalError()
}

pipeline = try! device?.makeComputePipelineState(function: function)

commandQueue = device?.makeCommandQueue()

CVMetalTextureCacheCreate(kCFAllocatorDefault, nil, device!, nil, &textureCache)

由於要讀取螢幕上顯示的畫面，需將MTKView.framebufferOnly屬性設定為false。

3. 編寫Gamma校正shader程式碼

inTexture表示攝像頭當前捕獲的畫面，outTexture表示處理後的資料，將會渲染到螢幕。

#include <metal_stdlib>
using namespace metal;

kernel void gamma_filter(
        texture2d<float, access::read> inTexture [[texture(0)]],
        texture2d<float, access::write> outTexture [[texture(1)]],
        uint2 gid [[thread_position_in_grid]])
{
    float4 inColor = inTexture.read(gid);
    const float4 outColor = float4(pow(inColor.rgb, float3(0.4/* gamma校正引數 */)), inColor.a);
    outTexture.write(outColor, gid);
}

4. 渲染Compute Shader處理後的紋理到螢幕

在MTKViewDelegate的draw(in view: MTKView)方法中繪製Compute Shader處理後的紋理到螢幕，參考程式碼如下。

guard let texture = imageTexture else {
    return
}
guard let drawable = view.currentDrawable else {
    return
}
guard let commandBuffer = commandQueue?.makeCommandBuffer() else {
    return
}

let encoder = commandBuffer.makeComputeCommandEncoder()
encoder.setComputePipelineState(pipeline!)
encoder.setTexture(texture, at: 0)
encoder.setTexture(drawable.texture, at: 1)

let threads = MTLSize(width: 16, height: 16, depth: 1)
let threadgroups = MTLSize(width: texture.width / threads.width,
                           height: texture.height / threads.height,
                           depth: 1)
encoder.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threads)
encoder.endEncoding()

commandBuffer.present(drawable)
commandBuffer.commit()

關鍵程式碼encoder.setTexture(drawable.texture, at: 1)指示compute shader將gamma校正結果寫到MTKView.currentDrawable.texture。

5. 讀取Metal渲染結果並生成UIImage

類似OpenGL ES的glReadPixels操作，需要注意大小端位元組序及UIKit與Metal紋理座標系的差異。由第4節渲染Compute Shader處理後的紋理到螢幕可知，MTKView.currentDrawable.texture是當前的渲染結果紋理，讀取Metal渲染結果問題就成了MTLTexture轉換成UIImage問題，可藉助Core Graphics介面實現，參考程式碼如下。

let image = currentDrawable?.texture.toUIImage()

為方便後續開發，給MTLTexture新增轉換成UIImage介面。

public extension MTLTexture {

    public func toUIImage() -> UIImage {
        let bytesPerPixel: Int = 4
        let imageByteCount = self.width * self.height * bytesPerPixel
        let bytesPerRow = self.width * bytesPerPixel
        var src = [UInt8](repeating: 0, count: Int(imageByteCount))

        let region = MTLRegionMake2D(0, 0, self.width, self.height)
        self.getBytes(&src, bytesPerRow: bytesPerRow, from: region, mipmapLevel: 0)
        let bitmapInfo = CGBitmapInfo(rawValue: (CGBitmapInfo.byteOrder32Little.rawValue | CGImageAlphaInfo.noneSkipFirst.rawValue))
        let colorSpace = CGColorSpaceCreateDeviceRGB()
        let bitsPerComponent = 8
        let context = CGContext(data: &src, width: self.width, height: self.height, bitsPerComponent: bitsPerComponent, bytesPerRow: bytesPerRow, space: colorSpace, bitmapInfo: bitmapInfo.rawValue);

        let dstImageFilter = context?.makeImage();

        return UIImage(cgImage: dstImageFilter!, scale: 0.0, orientation: UIImageOrientation.downMirrored) // 對於本文件，不需要downMirrored，因為第1節強制攝像頭輸出portrait方向影象
    }
}

6. 討論：compute shader合理的dispatchThreadgroups設定

第4節渲染Compute Shader處理後的紋理到螢幕簡單設定了dispatchThreadgroups，那麼合理的dispatchThreadgroups值應該是多少呢？可參考官方文件：Working with threads and threadgroups，參考設定程式碼如下。

let w = pipeline!.threadExecutionWidth
let h = pipeline!.maxTotalThreadsPerThreadgroup / w
let threadsPerThreadgroup = MTLSizeMake(w, h, 1)
let threadgroupsPerGrid = MTLSize(width: (texture.width + w - 1) / w,
                                  height: (texture.height + h - 1) / h,
                                  depth: 1)

使用上述程式碼，在iPhone 7p上計算1080p畫面，GPU耗時略有下降。