Introduce lut encoder (#49)

* introduce lut encoder

* change clamp to zero to clamp to edge

Author: s1ddok

Alloy/Encoders/LookUpTableEncoder.swift

@@ -0,0 +1,70 @@
+//
+//  LookUpTableEncoder.swift
+//  Alloy
+//
+//  Created by Andrey Volodin on 29.10.2019.
+//
+
+import Metal
+
+final public class LookUpTableEncoder {
+
+    // MARK: - Properties
+
+    public let pipelineState: MTLComputePipelineState
+    private let deviceSupportsNonuniformThreadgroups: Bool
+
+    // MARK: - Life Cycle
+
+    public convenience init(context: MTLContext) throws {
+        guard let library = context.shaderLibrary(for: type(of: self))
+        else { throw CommonErrors.metalInitializationFailed }
+        try self.init(library: library)
+    }
+
+    public init(library: MTLLibrary) throws {
+        self.deviceSupportsNonuniformThreadgroups = library.device.supports(feature: .nonUniformThreadgroups)
+        let constantValues = MTLFunctionConstantValues()
+        constantValues.set(self.deviceSupportsNonuniformThreadgroups,
+                           at: 0)
+        let functionName = type(of: self).functionName
+        self.pipelineState = try library.computePipelineState(function: functionName,
+                                                              constants: constantValues)
+    }
+
+    // MARK: - Encode
+
+    public func encode(sourceTexture: MTLTexture,
+                       outputTexture: MTLTexture,
+                       lut: MTLTexture,
+                       intensity: Float,
+                       in commandBuffer: MTLCommandBuffer) {
+        commandBuffer.compute { encoder in
+            encoder.label = "Look Up Table Encoder"
+            self.encode(sourceTexture: sourceTexture,
+                        outputTexture: outputTexture,
+                        lut: lut,
+                        intensity: intensity,
+                        using: encoder)
+        }
+    }
+
+    public func encode(sourceTexture: MTLTexture,
+                       outputTexture: MTLTexture,
+                       lut: MTLTexture,
+                       intensity: Float,
+                       using encoder: MTLComputeCommandEncoder) {
+        encoder.set(textures: [sourceTexture, outputTexture, lut])
+        encoder.set(intensity, at: 0)
+
+        if self.deviceSupportsNonuniformThreadgroups {
+            encoder.dispatch2d(state: pipelineState,
+                               exactly: outputTexture.size)
+        } else {
+            encoder.dispatch2d(state: pipelineState,
+                               covering: outputTexture.size)
+        }
+    }
+
+    public static let functionName = "lookUpTable"
+}

Alloy/Shaders/Shaders.metal

@@ -713,3 +713,30 @@ vertex float4 simpleVertex(constant float4* vertices [[ buffer(0) ]],
 fragment float4 plainColorFragment(constant float4& pointColor [[ buffer(0) ]]) {
     return pointColor;
 }
+
+// MARK: Look up table
+
+kernel void lookUpTable(texture2d<float, access::read> source [[ texture(0) ]],
+                        texture2d<float, access::write> destination [[ texture(1) ]],
+                        texture3d<float, access::sample> lut [[ texture(2) ]],
+                        constant float& intensity [[ buffer(0) ]],
+                        uint2 position [[thread_position_in_grid]]) {
+    const ushort2 textureSize = ushort2(destination.get_width(),
+                                        destination.get_height());
+    checkPosition(position, textureSize, deviceSupportsNonuniformThreadgroups);
+
+    constexpr sampler s(coord::normalized,
+                        address::clamp_to_edge,
+                        filter::linear);
+
+    // read original color
+    float4 sourceColor = source.read(position);
+
+    // use it to sample target color
+    sourceColor.rgb = mix(sourceColor.rgb,
+                          lut.sample(s, sourceColor.rgb).rgb,
+                          intensity);
+
+    // write it to destination texture
+    destination.write(sourceColor, position);
+}