introduce helpful kernels (#96)

Alloy/Encoders/StdMeanNormalization.swift

@@ -0,0 +1,88 @@
+import Metal
+import simd
+
+final public class StdMeanNormalization {
+
+    // MARK: - Propertires
+
+    public let pipelineState: MTLComputePipelineState
+    private let deviceSupportsNonuniformThreadgroups: Bool
+
+    // MARK: - Life Cycle
+
+    public convenience init(context: MTLContext) throws {
+        try self.init(library: context.library(for: Self.self))
+    }
+
+    public init(library: MTLLibrary) throws {
+        self.deviceSupportsNonuniformThreadgroups = library.device
+                                                           .supports(feature: .nonUniformThreadgroups)
+        let constantValues = MTLFunctionConstantValues()
+        constantValues.set(self.deviceSupportsNonuniformThreadgroups,
+                           at: 0)
+        self.pipelineState = try library.computePipelineState(function: Self.functionName,
+                                                              constants: constantValues)
+    }
+
+    // MARK: - Encode
+
+    public func callAsFunction(sourceTexture: MTLTexture,
+                               destinationTexture: MTLTexture,
+                               mean: SIMD3<Float>,
+                               std: SIMD3<Float>,
+                               in commandBuffer: MTLCommandBuffer) {
+        self.encode(sourceTexture: sourceTexture,
+                    destinationTexture: destinationTexture,
+                    mean: mean,
+                    std: std,
+                    in: commandBuffer)
+    }
+
+    public func callAsFunction(sourceTexture: MTLTexture,
+                               destinationTexture: MTLTexture,
+                               mean: SIMD3<Float>,
+                               std: SIMD3<Float>,
+                               using encoder: MTLComputeCommandEncoder) {
+        self.encode(sourceTexture: sourceTexture,
+                    destinationTexture: destinationTexture,
+                    mean: mean,
+                    std: std,
+                    using: encoder)
+    }
+
+    public func encode(sourceTexture: MTLTexture,
+                       destinationTexture: MTLTexture,
+                       mean: SIMD3<Float>,
+                       std: SIMD3<Float>,
+                       in commandBuffer: MTLCommandBuffer) {
+        commandBuffer.compute { encoder in
+            encoder.label = "Normalize Kernel"
+            self.encode(sourceTexture: sourceTexture,
+                        destinationTexture: destinationTexture,
+                        mean: mean,
+                        std: std,
+                        using: encoder)
+        }
+    }
+
+    public func encode(sourceTexture: MTLTexture,
+                       destinationTexture: MTLTexture,
+                       mean: SIMD3<Float>,
+                       std: SIMD3<Float>,
+                       using encoder: MTLComputeCommandEncoder) {
+        encoder.set(textures: [sourceTexture,
+                               destinationTexture])
+        encoder.set(mean, at: 0)
+        encoder.set(std, at: 1)
+
+        if self.deviceSupportsNonuniformThreadgroups {
+            encoder.dispatch2d(state: self.pipelineState,
+                               exactly: destinationTexture.size)
+        } else {
+            encoder.dispatch2d(state: self.pipelineState,
+                               covering: destinationTexture.size)
+        }
+    }
+
+    public static let functionName = "normalize"
+}

Alloy/Encoders/TextureAffineCrop.swift

@@ -1,4 +1,5 @@
 import Metal
+import simd
 
 final public class TextureAffineCrop {
 

Alloy/Encoders/TextureDivideByConstant.swift

@@ -0,0 +1,135 @@
+import Metal
+
+final public class TextureDivideByConstant {
+
+    // MARK: - Properties
+
+    public let pipelineState: MTLComputePipelineState
+    private let deviceSupportsNonuniformThreadgroups: Bool
+
+    // MARK: - Init
+
+    public convenience init(context: MTLContext,
+                            scalarType: MTLPixelFormat.ScalarType = .half) throws {
+        try self.init(library: context.library(for: Self.self),
+                      scalarType: scalarType)
+    }
+
+    public init(library: MTLLibrary,
+                scalarType: MTLPixelFormat.ScalarType = .half) throws {
+        self.deviceSupportsNonuniformThreadgroups = library.device
+                                                           .supports(feature: .nonUniformThreadgroups)
+
+        let constantValues = MTLFunctionConstantValues()
+        constantValues.set(self.deviceSupportsNonuniformThreadgroups,
+                           at: 0)
+        let functionName = Self.functionName + "_" + scalarType.rawValue
+        self.pipelineState = try library.computePipelineState(function: functionName,
+                                                              constants: constantValues)
+    }
+
+    // MARK: - Encode
+
+    public func callAsFunction(sourceTexture: MTLTexture,
+                               destinationTexture: MTLTexture,
+                               constant: SIMD4<Float>,
+                               in commandBuffer: MTLCommandBuffer) {
+        self.encode(sourceTexture: sourceTexture,
+                    destinationTexture: destinationTexture,
+                    constant: constant,
+                    in: commandBuffer)
+    }
+
+    public func callAsFunction(sourceTexture: MTLTexture,
+                               destinationTexture: MTLTexture,
+                               constant: SIMD4<Float>,
+                               using encoder: MTLComputeCommandEncoder) {
+        self.encode(sourceTexture: sourceTexture,
+                    destinationTexture: destinationTexture,
+                    constant: constant,
+                    using: encoder)
+    }
+
+    public func encode(sourceTexture: MTLTexture,
+                       destinationTexture: MTLTexture,
+                       constant: SIMD4<Float>,
+                       in commandBuffer: MTLCommandBuffer) {
+        commandBuffer.compute { encoder in
+            encoder.label = "Texture Divide by Constant"
+            self.encode(sourceTexture: sourceTexture,
+                        destinationTexture: destinationTexture,
+                        constant: constant,
+                        using: encoder)
+        }
+    }
+
+    public func encode(sourceTexture: MTLTexture,
+                       destinationTexture: MTLTexture,
+                       constant: SIMD4<Float>,
+                       using encoder: MTLComputeCommandEncoder) {
+        encoder.set(textures: [sourceTexture,
+                               destinationTexture])
+        encoder.set(constant, at: 0)
+
+        if self.deviceSupportsNonuniformThreadgroups {
+            encoder.dispatch2d(state: self.pipelineState,
+                               exactly: sourceTexture.size)
+        } else {
+            encoder.dispatch2d(state: self.pipelineState,
+                               covering: sourceTexture.size)
+        }
+    }
+
+    public func callAsFunction(sourceTexture: MTLTexture,
+                               destinationTexture: MTLTexture,
+                               constant: MTLBuffer,
+                               in commandBuffer: MTLCommandBuffer) {
+        self.encode(sourceTexture: sourceTexture,
+                    destinationTexture: destinationTexture,
+                    constant: constant,
+                    in: commandBuffer)
+    }
+
+    public func callAsFunction(sourceTexture: MTLTexture,
+                               destinationTexture: MTLTexture,
+                               constant: MTLBuffer,
+                               using encoder: MTLComputeCommandEncoder) {
+        self.encode(sourceTexture: sourceTexture,
+                    destinationTexture: destinationTexture,
+                    constant: constant,
+                    using: encoder)
+    }
+
+    public func encode(sourceTexture: MTLTexture,
+                       destinationTexture: MTLTexture,
+                       constant: MTLBuffer,
+                       in commandBuffer: MTLCommandBuffer) {
+        commandBuffer.compute { encoder in
+            encoder.label = "Texture Divide by Constant"
+            self.encode(sourceTexture: sourceTexture,
+                        destinationTexture: destinationTexture,
+                        constant: constant,
+                        using: encoder)
+        }
+    }
+
+    public func encode(sourceTexture: MTLTexture,
+                       destinationTexture: MTLTexture,
+                       constant: MTLBuffer,
+                       using encoder: MTLComputeCommandEncoder) {
+        encoder.set(textures: [sourceTexture,
+                               destinationTexture])
+        encoder.setBuffer(constant, offset: 0, index: 0)
+
+        if self.deviceSupportsNonuniformThreadgroups {
+            encoder.dispatch2d(state: self.pipelineState,
+                               exactly: sourceTexture.size)
+        } else {
+            encoder.dispatch2d(state: self.pipelineState,
+                               covering: sourceTexture.size)
+        }
+    }
+
+    public static let functionName = "divideByConstant"
+}
+

Alloy/Encoders/TextureNormalization.swift

@@ -0,0 +1,64 @@
+import Metal
+
+final public class TextureNormalization {
+
+    // MARK: - Properties
+
+    private let textureMax: TextureMax
+    private let textureDivide: TextureDivideByConstant
+    private let intermediateBuffer: MTLBuffer
+
+    // MARK: - Life Cycle
+
+    public convenience init(context: MTLContext) throws {
+        try self.init(library: context.library(for: Self.self))
+    }
+
+    public init(library: MTLLibrary) throws {
+        self.textureDivide = try .init(library: library)
+        self.textureMax = try .init(library: library)
+        self.intermediateBuffer = try library.device.buffer(for: SIMD4<Float>.self,
+                                                            options: .storageModePrivate)
+    }
+
+    // MARK: - Encode
+
+    public func callAsFunction(sourceTexture: MTLTexture,
+                               destinationTexture: MTLTexture,
+                               in commandBuffer: MTLCommandBuffer) {
+        self.encode(sourceTexture: sourceTexture,
+                    destinationTexture: destinationTexture,
+                    in: commandBuffer)
+    }
+
+    public func callAsFunction(sourceTexture: MTLTexture,
+                               destinationTexture: MTLTexture,
+                               using encoder: MTLComputeCommandEncoder) {
+        self.encode(sourceTexture: sourceTexture,
+                    destinationTexture: destinationTexture,
+                    using: encoder)
+    }
+
+    public func encode(sourceTexture: MTLTexture,
+                       destinationTexture: MTLTexture,
+                       in commandBuffer: MTLCommandBuffer) {
+        commandBuffer.compute { encoder in
+            encoder.label = "Texture Normalization"
+            self.encode(sourceTexture: sourceTexture,
+                        destinationTexture: destinationTexture,
+                        using: encoder)
+        }
+    }
+
+    public func encode(sourceTexture: MTLTexture,
+                       destinationTexture: MTLTexture,
+                       using encoder: MTLComputeCommandEncoder) {
+        self.textureMax(sourceTexture: sourceTexture,
+                        resultBuffer: self.intermediateBuffer,
+                        using: encoder)
+        self.textureDivide(sourceTexture: sourceTexture,
+                           destinationTexture: destinationTexture,
+                           constant: self.intermediateBuffer,
+                           using: encoder)
+    }
+}

Alloy/MTLTexture+Extensions.swift

@@ -1,6 +1,7 @@
 import Foundation
 import CoreGraphics
 import MetalKit
+import MetalPerformanceShaders
 import Accelerate
 
 public extension MTLTexture {
@@ -190,6 +191,19 @@ public extension MTLTexture {
         return matchingTexture
     }
 
+    func matchingTemporaryImage(commandBuffer: MTLCommandBuffer,
+                                usage: MTLTextureUsage? = nil) -> MPSTemporaryImage {
+        let matchingDescriptor = self.descriptor
+
+        if let u = usage {
+            matchingDescriptor.usage = u
+        }
+        // it has to be enforced for temporary image
+        matchingDescriptor.storageMode = .private
+
+        return MPSTemporaryImage(commandBuffer: commandBuffer, textureDescriptor: matchingDescriptor)
+    }
+
     func view(slice: Int,
               levels: Range<Int>? = nil) -> MTLTexture? {
         let sliceType: MTLTextureType

Alloy/Shaders/Shaders.metal

@@ -547,6 +547,38 @@ generateKernels(addConstant)
 #undef outerArguments
 #undef innerArguments
 
+template <typename T>
+void divideByConstant(texture2d<T, access::read> sourceTexture,
+                      texture2d<T, access::write> destinationTexture,
+                      constant float4& constantValue,
+                      const ushort2 position) {
+    const ushort2 textureSize = ushort2(sourceTexture.get_width(),
+                                        sourceTexture.get_height());
+    checkPosition(position, textureSize, deviceSupportsNonuniformThreadgroups);
+
+    auto sourceTextureValue = sourceTexture.read(position);
+    auto destinationTextureValue = sourceTextureValue / vec<T, 4>(constantValue);
+    destinationTexture.write(destinationTextureValue, position);
+}
+
+#define outerArguments(T)                                        \
+(texture2d<T, access::read> sourceTexture [[ texture(0) ]],      \
+texture2d<T, access::write> destinationTexture [[ texture(1) ]], \
+constant float4& constantValue [[ buffer(0) ]],                  \
+const ushort2 position [[ thread_position_in_grid ]])
+
+#define innerArguments \
+(sourceTexture,        \
+destinationTexture,    \
+constantValue,         \
+position)
+
+generateKernels(divideByConstant)
+
+#undef outerArguments
+#undef innerArguments
+
+
 // MARK: - Texture Mix
 
 kernel void textureMix(texture2d<float, access::read> sourceTextureOne [[ texture(0) ]],