import Accelerate import CoreGraphics import Metal import MetalKit import MetalPerformanceShaders import MPSX extension MTLCommandQueue { @discardableResult func sync(_ body: (MPSCommandBuffer) throws -> T) rethrows -> T { let commandBuffer = MPSCommandBuffer(from: self) let result = try autoreleasepool { try body(commandBuffer) } commandBuffer.commit() commandBuffer.waitUntilCompleted() return result } } extension CGImage { func jpeg(compressionQuality: Float = 1.0) -> Data? { guard let mutableData = CFDataCreateMutable(kCFAllocatorDefault, 0), let destination = CGImageDestinationCreateWithData(mutableData, "public.jpeg" as CFString, 1, nil) else { return nil } let properties: [CFString: Any] = [ kCGImageDestinationLossyCompressionQuality: compressionQuality as CFNumber, ] CGImageDestinationAddImage(destination, self, properties as CFDictionary) guard CGImageDestinationFinalize(destination) else { return nil } return mutableData as Data } } extension MTLTexture { func cgImage(colorSpace: CGColorSpace? = nil) -> CGImage? { switch pixelFormat { case .r8Unorm: let rowBytes = width let length = rowBytes * height let bytes = UnsafeMutableRawPointer.allocate( byteCount: length, alignment: MemoryLayout.alignment ) defer { bytes.deallocate() } getBytes( bytes, bytesPerRow: rowBytes, from: MTLRegionMake2D(0, 0, width, height), mipmapLevel: 0 ) guard let data = CFDataCreate( nil, bytes.assumingMemoryBound(to: UInt8.self), length ), let dataProvider = CGDataProvider(data: data), let cgImage = CGImage( width: width, height: height, bitsPerComponent: 8, bitsPerPixel: 8, bytesPerRow: rowBytes, space: colorSpace ?? CGColorSpaceCreateDeviceGray(), bitmapInfo: .init(rawValue: CGImageAlphaInfo.none.rawValue), provider: dataProvider, decode: nil, shouldInterpolate: true, intent: .defaultIntent ) else { return nil } return cgImage case .rgba8Unorm, .rgba8Unorm_srgb: let rowBytes = width * 4 let length = rowBytes * height let bytes = UnsafeMutableRawPointer.allocate( byteCount: length, alignment: MemoryLayout.alignment ) defer { bytes.deallocate() } getBytes( bytes, bytesPerRow: rowBytes, from: MTLRegionMake2D(0, 0, width, height), mipmapLevel: 0 ) guard let data = CFDataCreate( nil, bytes.assumingMemoryBound(to: UInt8.self), length ), let dataProvider = CGDataProvider(data: data), let cgImage = CGImage( width: width, height: height, bitsPerComponent: 8, bitsPerPixel: 32, bytesPerRow: rowBytes, space: colorSpace ?? CGColorSpace(name: CGColorSpace.displayP3)!, bitmapInfo: .init(rawValue: CGImageAlphaInfo.premultipliedLast.rawValue), provider: dataProvider, decode: nil, shouldInterpolate: false, intent: .defaultIntent ) else { return nil } return cgImage default: return nil } } } final class GPU { // MARK: Lifecycle init(commandQueue: MTLCommandQueue) { self.commandQueue = commandQueue textureLoader = .init(device: commandQueue.device) } // MARK: Internal static let `default` = GPU(commandQueue: MTLCreateSystemDefaultDevice()!.makeCommandQueue()!) let commandQueue: MTLCommandQueue let textureLoader: MTKTextureLoader var device: MTLDevice { commandQueue.device } } func compare(texture: MTLTexture, with reference: MTLTexture, treshold: Float = 1e-3) -> Bool { let graph = MPSCompiledGraph( device: GPU.default.device, options: .init(runtimeTypeInference: true) ) { graph in let x = graph.imagePlaceholder( dataType: .float32, height: reference.height, width: reference.width, channels: -1, name: "X" ) let y = graph.imagePlaceholder( dataType: .float32, height: reference.height, width: reference.width, channels: -1, name: "Y" ) let z = (x.mean(axes: [3]) - y.mean(axes: [3])).pow(2).mean(axes: [1, 2]) return ["Z": z] } let result = GPU.default.commandQueue.sync { graph([ "X": .texture(texture), "Y": .texture(reference), ], in: $0).synchronizedNDArray(in: $0) }.floats return result[0] < treshold } let testResourcesPath = "TestResources" func data(bundlePath: String) throws -> Data { try Data(contentsOf: Bundle.module.url(forResource: bundlePath, withExtension: nil)!) } func texture(bundlePath: String) async throws -> MTLTexture { try await GPU.default.textureLoader.newTexture( URL: .init( fileURLWithPath: Bundle.module.path(forResource: bundlePath, ofType: nil)! ), options: [.SRGB: false] ) } func data(arg: Int) throws -> Data { try Data(contentsOf: URL(fileURLWithPath: CommandLine.arguments[arg])) } func inputTexture(arg: Int) async throws -> MTLTexture { try await GPU.default.textureLoader.newTexture( URL: .init( fileURLWithPath: CommandLine.arguments[arg] ), options: [.SRGB: false] ) } func save(texture: MTLTexture, arg: Int) throws { let image = texture.cgImage() try image?.jpeg(compressionQuality: 0.85)?.write( to: .init( fileURLWithPath: CommandLine.arguments[arg] ), options: .atomic ) }