diff --git a/Alloy/Encoders/StdMeanNormalization.swift b/Alloy/Encoders/StdMeanNormalization.swift new file mode 100644 index 0000000..0aee8e5 --- /dev/null +++ b/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, + std: SIMD3, + in commandBuffer: MTLCommandBuffer) { + self.encode(sourceTexture: sourceTexture, + destinationTexture: destinationTexture, + mean: mean, + std: std, + in: commandBuffer) + } + + public func callAsFunction(sourceTexture: MTLTexture, + destinationTexture: MTLTexture, + mean: SIMD3, + std: SIMD3, + using encoder: MTLComputeCommandEncoder) { + self.encode(sourceTexture: sourceTexture, + destinationTexture: destinationTexture, + mean: mean, + std: std, + using: encoder) + } + + public func encode(sourceTexture: MTLTexture, + destinationTexture: MTLTexture, + mean: SIMD3, + std: SIMD3, + 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, + std: SIMD3, + 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" +} diff --git a/Alloy/Encoders/TextureAffineCrop.swift b/Alloy/Encoders/TextureAffineCrop.swift index 28ef21e..a9dc9ae 100644 --- a/Alloy/Encoders/TextureAffineCrop.swift +++ b/Alloy/Encoders/TextureAffineCrop.swift @@ -1,4 +1,5 @@ import Metal +import simd final public class TextureAffineCrop { diff --git a/Alloy/Encoders/TextureDivideByConstant.swift b/Alloy/Encoders/TextureDivideByConstant.swift new file mode 100644 index 0000000..b5a3510 --- /dev/null +++ b/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, + in commandBuffer: MTLCommandBuffer) { + self.encode(sourceTexture: sourceTexture, + destinationTexture: destinationTexture, + constant: constant, + in: commandBuffer) + } + + public func callAsFunction(sourceTexture: MTLTexture, + destinationTexture: MTLTexture, + constant: SIMD4, + using encoder: MTLComputeCommandEncoder) { + self.encode(sourceTexture: sourceTexture, + destinationTexture: destinationTexture, + constant: constant, + using: encoder) + } + + public func encode(sourceTexture: MTLTexture, + destinationTexture: MTLTexture, + constant: SIMD4, + 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, + 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" +} + diff --git a/Alloy/Encoders/TextureNormalization.swift b/Alloy/Encoders/TextureNormalization.swift new file mode 100644 index 0000000..6de6f6e --- /dev/null +++ b/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.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) + } +} diff --git a/Alloy/MTLTexture+Extensions.swift b/Alloy/MTLTexture+Extensions.swift index fc6556c..40cfaaa 100644 --- a/Alloy/MTLTexture+Extensions.swift +++ b/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? = nil) -> MTLTexture? { let sliceType: MTLTextureType diff --git a/Alloy/Shaders/Shaders.metal b/Alloy/Shaders/Shaders.metal index 21e080b..1eea83b 100644 --- a/Alloy/Shaders/Shaders.metal +++ b/Alloy/Shaders/Shaders.metal @@ -547,6 +547,38 @@ generateKernels(addConstant) #undef outerArguments #undef innerArguments +template +void divideByConstant(texture2d sourceTexture, + texture2d 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(constantValue); + destinationTexture.write(destinationTextureValue, position); +} + +#define outerArguments(T) \ +(texture2d sourceTexture [[ texture(0) ]], \ +texture2d 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 sourceTextureOne [[ texture(0) ]],