Metal学习(6) - 在计算函数中处理纹理

通过计算函数先把输入纹理进行处理，得到新的纹理后，用新纹理进行渲染。

ShaderType

#include <simd/simd.h>

typedef enum ShaderVertexInputIndex
{
    ShaderVertexInputIndexVertices     = 0,
    ShaderVertexInputIndexViewportSize = 1,
} ShaderVertexInputIndex;

/// 纹理索引
typedef enum ShaderTextureIndex
{
    ShaderTextureIndexInput  = 0,
    ShaderTextureIndexOutput = 1,
} ShaderTextureIndex;

typedef struct
{
    vector_float2 position;
    vector_float2 textureCoordinate;
} ShaderVertex;

Shader.metal

#include <metal_stdlib>
using namespace metal;

#include "ShaderType.h"

struct RasterizerData
{
    float4 position [[position]];
    float2 textureCoordinate;
};

vertex RasterizerData vertexShader(const uint vertexID [[vertex_id]],
                                   constant ShaderVertex *vertices [[buffer(ShaderVertexInputIndexVertices)]])
{
    RasterizerData out;
    out.position = vector_float4(vertices[vertexID].position.x, vertices[vertexID].position.y, 0.0, 1.0);
    out.textureCoordinate = vertices[vertexID].textureCoordinate;
    return out;
}

fragment float4 fragmentShader(RasterizerData in [[stage_in]],
                               texture2d<half> colorTexture [[texture(ShaderTextureIndexOutput)]])
{
    constexpr sampler textureSampler (mag_filter::linear, min_filter::linear);
    const half4 colorSample = colorTexture.sample(textureSampler, in.textureCoordinate);
    return float4(colorSample);
}

constant half3 kRec709Luma = half3(0.2126, 0.7152, 0.0722);
// 从输入纹理，进行灰度计算，写入输出纹理
kernel void grayscaleKernel(texture2d<half, access::read>  inTexture  [[texture(ShaderTextureIndexInput)]],
                            texture2d<half, access::write> outTexture [[texture(ShaderTextureIndexOutput)]],
                            uint2                          gid        [[thread_position_in_grid]])
{
    // 判断是否超出纹理了
    if((gid.x >= outTexture.get_width()) || (gid.y >= outTexture.get_height())) {
        return;
    }
    half4 inColor = inTexture.read(gid);
    half  gray = dot(inColor.rgb, kRec709Luma);
    outTexture.write(half4(gray, gray, gray, 1.0), gid);
}

MetalRender.swift

import UIKit
import MetalKit
import simd

class MetalRender: NSObject {

    // 向设备传递命令的命令队列
    private var commandQueue: MTLCommandQueue?
    // 计算管线状态
    private var computePipelineState: MTLComputePipelineState?
    // 渲染管线状态
    private var renderPipelineState: MTLRenderPipelineState?
    private var inputTexture: MTLTexture?
    private var outputTexture: MTLTexture?
    // 线程组的大小
    private let threadgroupSize = MTLSize(width: 16, height: 16, depth: 1)
    // 线程组数量
    private var threadgroupCount: MTLSize?
    // 渲染顶点
    private let triangleVertices: [ShaderVertex] = {
        let vertex0 = ShaderVertex(position: vector_float2(x: -1, y: -1), textureCoordinate: vector_float2(x: 0, y: 0))
        let vertex1 = ShaderVertex(position: vector_float2(-1, 1), textureCoordinate: vector_float2(0, 1))
        let vertex2 = ShaderVertex(position: vector_float2(1, -1), textureCoordinate: vector_float2(1, 0))
        let vertex3 = ShaderVertex(position: vector_float2(1, 1), textureCoordinate: vector_float2(1, 1))
        let array: [ShaderVertex] = [vertex0, vertex1, vertex2, vertex3]
        return array
    }()

    private override init() {
        super.init()
    }
    convenience init(_ view: MTKView) {
        self.init()

        let device = view.device

        let defaultLibrary = device?.makeDefaultLibrary()
        let kernelFunction = defaultLibrary?.makeFunction(name: "grayscaleKernel")
        let vertexFunction = defaultLibrary?.makeFunction(name: "vertexShader")
        let fragmentFunction = defaultLibrary?.makeFunction(name: "fragmentShader")

        let pipelineStateDescriptor = MTLRenderPipelineDescriptor()
        pipelineStateDescriptor.label = "Simple Pipeline"
        pipelineStateDescriptor.vertexFunction = vertexFunction
        pipelineStateDescriptor.fragmentFunction = fragmentFunction
        pipelineStateDescriptor.colorAttachments[0].pixelFormat = view.colorPixelFormat
        do {
            // 创建计算管线
            computePipelineState = try device?.makeComputePipelineState(function: kernelFunction!)
            // 创建渲染管线
            renderPipelineState = try device?.makeRenderPipelineState(descriptor: pipelineStateDescriptor)
        }catch {
            print(error)
        }
        // 命令队列
        commandQueue = device?.makeCommandQueue()

        // 通过图片创建inputTexture纹理
        guard let image = flipImage(UIImage(named: "02")), let cgImage = image.cgImage else { return }
        let textureLoader = MTKTextureLoader(device: device!)
        inputTexture = try? textureLoader.newTexture(cgImage: cgImage, options: [MTKTextureLoader.Option.SRGB : false])

        // 根据inputTexture纹理的信息，创建outputTexture纹理，这里必须先创建纹理，才能把计算结果写入纹理
        let descriptor = MTLTextureDescriptor()
        descriptor.textureType = inputTexture!.textureType
        descriptor.pixelFormat = inputTexture!.pixelFormat
        descriptor.width = inputTexture!.width
        descriptor.height = inputTexture!.height
        // 因为outputTexture是computePipeline的输出，也是renderPipeline的输入，所以同时支持读和写
        descriptor.usage = [.shaderWrite, .shaderRead]
        outputTexture = device!.makeTexture(descriptor: descriptor)
        
        // 因为本次使用的是均匀线程组，所以分配的实际区域要大于等于纹理大小
        threadgroupCount = MTLSize(width: (inputTexture!.width+threadgroupSize.width-1)/threadgroupSize.width, height: (inputTexture!.height+threadgroupSize.height-1)/threadgroupSize.height, depth: 1)
    }
}

// MARK: - private

private extension MetalRender {
    /// 翻转上下颠倒的图片
    func flipImage(_ image: UIImage?) -> UIImage? {
        guard let image = image else  { return nil }
        UIGraphicsBeginImageContextWithOptions(image.size, false, UIScreen.main.scale)
        let ctx = UIGraphicsGetCurrentContext()
        ctx?.translateBy(x: 0, y: image.size.height)
        ctx?.scaleBy(x: 1, y: -1)
        image.draw(in: CGRect(origin: .zero, size: image.size))
        let reslut = UIGraphicsGetImageFromCurrentImageContext()
        UIGraphicsEndImageContext()
        return reslut
    }
}


// MARK: - MTKViewDelegate

extension MetalRender: MTKViewDelegate {
    // 每当视图改变方向或调整大小时调用
    func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {
    }
    // 当视图需要渲染帧时调
    func draw(in view: MTKView) {

        guard computePipelineState != nil else { return }
        guard renderPipelineState != nil else { return }

        // 创建命令缓冲区
        let commandBuffer = commandQueue?.makeCommandBuffer()
        commandBuffer?.label = "MyCommand"

        // 创建计算命令编码器
        let computeEncoder = commandBuffer?.makeComputeCommandEncoder()
        computeEncoder?.label = "MyComputeEncoder"
        computeEncoder?.setComputePipelineState(computePipelineState!)
        computeEncoder?.setTexture(inputTexture, index: Int(ShaderTextureIndexInput.rawValue))
        computeEncoder?.setTexture(outputTexture, index: Int(ShaderTextureIndexOutput.rawValue))
        // 采用均匀线程组
        computeEncoder?.dispatchThreadgroups(threadgroupCount!, threadsPerThreadgroup: threadgroupSize)
        computeEncoder?.endEncoding()

        if let renderPassDescriptor = view.currentRenderPassDescriptor {
            // 创建渲染命令编码器
            let renderEncoder = commandBuffer?.makeRenderCommandEncoder(descriptor: renderPassDescriptor)
            renderEncoder?.label = "MyRenderEncoder"
            renderEncoder?.setViewport(MTLViewport(originX: 0, originY: 0, width: Double(view.drawableSize.width), height: Double(view.drawableSize.height), znear: -1, zfar: 1))
            renderEncoder?.setRenderPipelineState(renderPipelineState!)
            renderEncoder?.setVertexBytes(triangleVertices, length: triangleVertices.count*MemoryLayout<ShaderVertex>.size, index: Int(ShaderVertexInputIndexVertices.rawValue))
            // 把outputTexture作为输入的纹理传入
            renderEncoder?.setFragmentTexture(outputTexture, index: Int(ShaderTextureIndexOutput.rawValue))
            renderEncoder?.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4)
            renderEncoder?.endEncoding()

            if let drawable = view.currentDrawable {
                commandBuffer?.present(drawable)
            }
            
            commandBuffer?.addCompletedHandler({ buffer in
                let outputTextureCIImage = CIImage(mtlTexture: self.outputTexture!)
                let outputTextureImage = UIImage(ciImage: outputTextureCIImage!)
                print("可以得到处理后的纹理图片")
            })
        }
        commandBuffer?.commit()
    }
}

需要注意的是，计算函数的输出纹理的usage，必须包含.shaderWrite