diff --git a/Sources/mtlswift/AST/ASTNode.swift b/Sources/mtlswift/AST/ASTNode.swift index 2e5ecb5..6e8c3f4 100644 --- a/Sources/mtlswift/AST/ASTNode.swift +++ b/Sources/mtlswift/AST/ASTNode.swift @@ -21,7 +21,7 @@ public class ASTNode { public init(parsingString inputString: String) throws { let scanner = StringScanner(string: inputString) - + guard let prefix = scanner.readWord() else { throw Errors.parsingError } diff --git a/Sources/mtlswift/AST/ASTShader.swift b/Sources/mtlswift/AST/ASTShader.swift index da39f17..8bc0b4a 100644 --- a/Sources/mtlswift/AST/ASTShader.swift +++ b/Sources/mtlswift/AST/ASTShader.swift @@ -58,14 +58,18 @@ public struct ASTShader { var swiftNameLookup: [String: String] = [:] var swiftTypeLookup: [String: String] = [:] + var inPlaceTextureNameMappings = [MTLKernelEncoder.InPlaceTextureNameMapping]() for declaration in self.customDeclarations { - if case .swiftParameterName(let oldName, let newName) = declaration { + if case let .swiftParameterName(oldName, newName) = declaration { swiftNameLookup[oldName] = newName } - if case .swiftParameterType(let parameter, let type) = declaration { + if case let .swiftParameterType(parameter, type) = declaration { swiftTypeLookup[parameter] = type } + if case let .inPlaceTexture(source, destination, inPlace) = declaration { + inPlaceTextureNameMappings.append(.init(source: source, destination: destination, inPlace: inPlace)) + } } // individual access levels are not supported for now @@ -92,7 +96,7 @@ public struct ASTShader { print("WARNING: Swift Types are not available for texture parameters, ignoring \(type)") type = "MTLTexture" } - + return MTLKernelEncoder.Parameter(name: name, swiftTypeName: type, kind: .texture, @@ -172,7 +176,8 @@ public struct ASTShader { encodingVariants: [MTLKernelEncoder.EncodingVariant(dispatchType: type, threadgroupSize: size)], usedConstants: constants, branchingConstant: branchingConstant, - threadgroupMemoryCalculations: threadgroupMemoryCalculatiosn) + threadgroupMemoryCalculations: threadgroupMemoryCalculatiosn, + inPlaceTextureNameMappings: inPlaceTextureNameMappings) } return nil diff --git a/Sources/mtlswift/AST/CustomDeclaration.swift b/Sources/mtlswift/AST/CustomDeclaration.swift index 27462e9..5e43130 100644 --- a/Sources/mtlswift/AST/CustomDeclaration.swift +++ b/Sources/mtlswift/AST/CustomDeclaration.swift @@ -4,7 +4,8 @@ public enum CustomDeclaration { let scanner = StringScanner(string: rawString) if scanner.skip(exact: CustomDeclaration.swiftNameDeclaration) { - guard let name = scanner.readWord() else { + guard let name = scanner.readWord() + else { print("ERROR: Failed to parse \(CustomDeclaration.swiftNameDeclaration), skipping: \(scanner.leftString)") return nil } @@ -12,6 +13,21 @@ public enum CustomDeclaration { // TODO: Check it is valid Swift identifier self = .swiftName(name: name) return + } else if scanner.skip(exact: CustomDeclaration.inPlaceTexture) { + guard let source = scanner.readWord(), + scanner.skip(exact: ":"), + let destination = scanner.readWord(), + scanner.skip(exact: ":"), + let inPlace = scanner.readWord() + else { + print("ERROR: Failed to parse \(CustomDeclaration.inPlaceTexture), skipping: \(scanner.leftString)") + return nil + } + + self = .inPlaceTexture(source: source, + destination: destination, + inPlace: inPlace) + return } else if scanner.skip(exact: CustomDeclaration.swiftParameterNameDeclaration) { guard let oldName = scanner.readWord(), @@ -186,6 +202,9 @@ public enum CustomDeclaration { public static let accessLevelDeclaration = "accessLevel:" case accessLevel(level: AccessLevel) + + public static let inPlaceTexture = "inPlaceTexture:" + case inPlaceTexture(source: String, destination: String, inPlace: String) public static let swiftNameDeclaration = "swiftName:" case swiftName(name: String) diff --git a/Sources/mtlswift/Encoder Generator/MTLKernelEncoder.swift b/Sources/mtlswift/Encoder Generator/MTLKernelEncoder.swift index 8ace7b0..24269f7 100644 --- a/Sources/mtlswift/Encoder Generator/MTLKernelEncoder.swift +++ b/Sources/mtlswift/Encoder Generator/MTLKernelEncoder.swift @@ -1,6 +1,6 @@ public struct MTLKernelEncoder { - public struct Parameter { + public struct Parameter: Hashable { public enum Kind { case texture, buffer, sampler, threadgroupMemory } @@ -12,6 +12,12 @@ public struct MTLKernelEncoder { public var defaultValueString: String? = nil } + + public struct InPlaceTextureNameMapping { + public let source: String + public let destination: String + public let inPlace: String + } public enum ThreadgroupMemoryLengthCalculation { case total(index: Int, bytes: Int) @@ -32,6 +38,7 @@ public struct MTLKernelEncoder { public var usedConstants: [ASTFunctionConstant] public var branchingConstant: ASTFunctionConstant? public var threadgroupMemoryCalculations: [ThreadgroupMemoryLengthCalculation] + public var inPlaceTextureNameMappings: [InPlaceTextureNameMapping] public var shaderString: String { let sourceBuilder = SourceStringBuilder() @@ -45,19 +52,34 @@ public struct MTLKernelEncoder { sourceBuilder.add(line: "\(self.accessLevel.rawValue) let \(bc.name): \(bc.type.swiftTypeDelcaration)") } + let containsInPlaceTextures = !self.inPlaceTextureNameMappings.isEmpty + sourceBuilder.blankLine() sourceBuilder.add(line: "\(self.accessLevel.rawValue) let pipelineState: MTLComputePipelineState") + if containsInPlaceTextures { + sourceBuilder.add(line: "\(self.accessLevel.rawValue) let textureCopy: TextureCopy") + } sourceBuilder.blankLine() if self.usedConstants.isEmpty && self.branchingConstant == nil { // MARK: Generate inits - sourceBuilder.add(line: "\(self.accessLevel.rawValue) init(library: MTLLibrary) throws {") + if containsInPlaceTextures { + sourceBuilder.add(line: "\(self.accessLevel.rawValue) init(library: MTLLibrary, textureCopy: TextureCopy) throws {") + } else { + sourceBuilder.add(line: "\(self.accessLevel.rawValue) init(library: MTLLibrary) throws {") + } sourceBuilder.pushLevel() sourceBuilder.add(line: "self.pipelineState = try library.computePipelineState(function: \"\(self.shaderName)\")") } else { let parameterString = ", " + self.usedConstants.map { "\($0.name): \($0.type.swiftTypeDelcaration)" }.joined(separator: ", ") - sourceBuilder.add(line: "\(self.accessLevel.rawValue) init(library: MTLLibrary\(self.usedConstants.isEmpty ? "" : parameterString)) throws {") + + if containsInPlaceTextures { + sourceBuilder.add(line: "\(self.accessLevel.rawValue) init(library: MTLLibrary, textureCopy: TextureCopy\(self.usedConstants.isEmpty ? "" : parameterString)) throws {") + } else { + sourceBuilder.add(line: "\(self.accessLevel.rawValue) init(library: MTLLibrary\(self.usedConstants.isEmpty ? "" : parameterString)) throws {") + } + sourceBuilder.pushLevel() sourceBuilder.add(line: "let constantValues = MTLFunctionConstantValues()") @@ -76,6 +98,10 @@ public struct MTLKernelEncoder { sourceBuilder.add(line: "self.pipelineState = try library.computePipelineState(function: \"\(self.shaderName)\", constants: constantValues)") } + if containsInPlaceTextures { + sourceBuilder.add(line: "self.textureCopy = textureCopy") + } + // MARK: Balancing for init sourceBuilder.popLevel() sourceBuilder.add(line: "}") @@ -113,11 +139,48 @@ public struct MTLKernelEncoder { sourceBuilder.add(line: "\(self.accessLevel.rawValue) func encode(\(gridSizeParameterString)\(threadgroupParameterString)using encoder: MTLComputeCommandEncoder) {") } else { var parameterString = "" + var filteredParameterString = "" + + let inPlaceParameters = self.inPlaceTextureNameMappings.map { mapping -> Parameter in + return .init(name: mapping.inPlace, + swiftTypeName: "MTLTexture", + kind: .texture, + index: -1) + } + let sourceParameters = self.inPlaceTextureNameMappings.map { mapping -> Parameter in + let index = self.parameters.first(where: { $0.name == mapping.source })?.index ?? 0 + return .init(name: mapping.source, + swiftTypeName: "MTLTexture", + kind: .texture, + index: index) + } + let destinationParameters = self.inPlaceTextureNameMappings.map { mapping -> Parameter in + let index = self.parameters.first(where: { $0.name == mapping.destination })?.index ?? 0 + return .init(name: mapping.destination, + swiftTypeName: "MTLTexture", + kind: .texture, + index: index) + } + let filteredParameters = containsInPlaceTextures + ? inPlaceParameters + Set(self.parameters).subtracting(Set(sourceParameters)) + .subtracting(Set(destinationParameters)) + .map { $0 } + : self.parameters + for parameter in self.parameters { parameterString += "\(parameter.name): \(parameter.swiftTypeName), " } + + for filteredParameter in filteredParameters { + filteredParameterString += "\(filteredParameter.name): \(filteredParameter.swiftTypeName), " + } + + if containsInPlaceTextures { + print(2) + } var parametersBodyString = "" + var filteredParametersBodyString = "" let gridSizeValueString = gridSizeParameterString.isEmpty ? "" : ", gridSize: gridSize" let threadgroupSizeValueString = threadgroupParameterString.isEmpty ? "" : ", threadgroupSize: threadgroupSize" for parameterIndex in 0 ..< self.parameters.count { @@ -125,24 +188,57 @@ public struct MTLKernelEncoder { let parameterSeparator = parameterIndex < self.parameters.count - 1 ? ", " : "" parametersBodyString += parameterName + ": " + parameterName + parameterSeparator } + for parameterIndex in 0 ..< filteredParameters.count { + let parameterName = filteredParameters[parameterIndex].name + let parameterSeparator = parameterIndex < filteredParameters.count - 1 ? ", " : "" + filteredParametersBodyString += parameterName + ": " + parameterName + parameterSeparator + } // Call as function in commandBuffer - sourceBuilder.add(line: "\(self.accessLevel.rawValue) func callAsFunction(\(parameterString)\(gridSizeParameterString)\(threadgroupParameterString)in commandBuffer: MTLCommandBuffer) {") + sourceBuilder.add(line: "\(self.accessLevel.rawValue) func callAsFunction(\(filteredParameterString)\(gridSizeParameterString)\(threadgroupParameterString)in commandBuffer: MTLCommandBuffer) {") sourceBuilder.pushLevel() - sourceBuilder.add(line: "self.encode(\(parametersBodyString)\(gridSizeValueString)\(threadgroupSizeValueString), in: commandBuffer)") + sourceBuilder.add(line: "self.encode(\(filteredParametersBodyString)\(gridSizeValueString)\(threadgroupSizeValueString), in: commandBuffer)") sourceBuilder.popLevel() sourceBuilder.add(line: "}") - // Call as function using encoder - sourceBuilder.add(line: "\(self.accessLevel.rawValue) func callAsFunction(\(parameterString)\(gridSizeParameterString)\(threadgroupParameterString)using encoder: MTLComputeCommandEncoder) {") - sourceBuilder.pushLevel() - sourceBuilder.add(line: "self.encode(\(parametersBodyString)\(gridSizeValueString)\(threadgroupSizeValueString), using: encoder)") - sourceBuilder.popLevel() - sourceBuilder.add(line: "}") + if !containsInPlaceTextures { + // Call as function using encoder + sourceBuilder.add(line: "\(self.accessLevel.rawValue) func callAsFunction(\(filteredParameterString)\(gridSizeParameterString)\(threadgroupParameterString)using encoder: MTLComputeCommandEncoder) {") + sourceBuilder.pushLevel() + sourceBuilder.add(line: "self.encode(\(filteredParametersBodyString)\(gridSizeValueString)\(threadgroupSizeValueString), using: encoder)") + sourceBuilder.popLevel() + sourceBuilder.add(line: "}") + } // Encode in commandBuffer - sourceBuilder.add(line: "\(self.accessLevel.rawValue) func encode(\(parameterString)\(gridSizeParameterString)\(threadgroupParameterString)in commandBuffer: MTLCommandBuffer) {") + sourceBuilder.add(line: "\(self.accessLevel.rawValue) func encode(\(filteredParameterString)\(gridSizeParameterString)\(threadgroupParameterString)in commandBuffer: MTLCommandBuffer) {") sourceBuilder.pushLevel() + + if containsInPlaceTextures { + for inPlaceTextureNameMapping in self.inPlaceTextureNameMappings { + let sourceName = inPlaceTextureNameMapping.source + let destinationName = inPlaceTextureNameMapping.destination + let inPlaceName = inPlaceTextureNameMapping.inPlace + let imageCopyName = "\(sourceName)CopyImage" + let originalTextureName = "\(sourceName)OriginalTexture" + + sourceBuilder.add(line: "var \(sourceName) = \(inPlaceName)") + sourceBuilder.add(line: "let \(destinationName) = \(inPlaceName)") + sourceBuilder.add(line: "if !self.pipelineState.device.supports(feature: .readWriteTextures(\(sourceName).pixelFormat)) {") + sourceBuilder.pushLevel() + + sourceBuilder.add(line: "let \(originalTextureName) = \(sourceName)") + sourceBuilder.add(line: "let \(imageCopyName) = \(sourceName).matchingTemporaryImage(commandBuffer: commandBuffer)") + sourceBuilder.add(line: "defer { \(imageCopyName).readCount = .zero }") + sourceBuilder.add(line: "\(sourceName) = \(imageCopyName).texture") + sourceBuilder.add(line: "self.textureCopy(source: \(originalTextureName), destination: \(sourceName), in: commandBuffer)") + sourceBuilder.add(line: "}") + sourceBuilder.popLevel() + + sourceBuilder.blankLine() + } + } + sourceBuilder.add(line: "commandBuffer.compute { encoder in") sourceBuilder.pushLevel() sourceBuilder.add(line: "encoder.label = \"\(self.swiftName)\"") @@ -153,7 +249,8 @@ public struct MTLKernelEncoder { sourceBuilder.add(line: "}") // Ecode using encoder - sourceBuilder.add(line: "\(self.accessLevel.rawValue) func encode(\(parameterString)\(gridSizeParameterString)\(threadgroupParameterString)using encoder: MTLComputeCommandEncoder) {") + let accessLevel = containsInPlaceTextures ? "private" : "\(self.accessLevel.rawValue)" + sourceBuilder.add(line: "\(accessLevel) func encode(\(parameterString)\(gridSizeParameterString)\(threadgroupParameterString)using encoder: MTLComputeCommandEncoder) {") } sourceBuilder.pushLevel() sourceBuilder.add(line: threadgroupVariableString) @@ -201,7 +298,7 @@ public struct MTLKernelEncoder { case .even(parameters: .over(let argument)): if let targetParameter = self.parameters.first(where: { $0.name == argument }), - targetParameter.kind == .texture { + (targetParameter.kind == .texture) { sourceBuilder.add(line: "encoder.dispatch2d(state: self.pipelineState, covering: \(targetParameter.name).size\(threadgroupExpressionString))") } else { fatalError("Could not generate dispatching over parameter \(argument)") @@ -216,7 +313,7 @@ public struct MTLKernelEncoder { case .exact(parameters: .over(let argument)): if let targetParameter = self.parameters.first(where: { $0.name == argument }), - targetParameter.kind == .texture { + (targetParameter.kind == .texture) { sourceBuilder.add(line: "encoder.dispatch2d(state: self.pipelineState, exactly: \(targetParameter.name).size\(threadgroupExpressionString))") } else { print("Could not generate dispatching over parameter \(argument)") @@ -231,7 +328,7 @@ public struct MTLKernelEncoder { sourceBuilder.add(line: "if self.\(bc.name) { encoder.dispatch2d(state: self.pipelineState, exactly: \(self.swiftName).gridSize\(idx)\(threadgroupExpressionString)) } else { encoder.dispatch2d(state: self.pipelineState, covering: \(self.swiftName).gridSize\(idx)\(threadgroupExpressionString)) }") case .optimal(_, parameters: .over(let argument)): if let targetParameter = self.parameters.first(where: { $0.name == argument }), - targetParameter.kind == .texture { + (targetParameter.kind == .texture) { let bc = self.branchingConstant! sourceBuilder.add(line: "if self.\(bc.name) { encoder.dispatch2d(state: self.pipelineState, exactly: \(targetParameter.name).size\(threadgroupExpressionString)) } else { encoder.dispatch2d(state: self.pipelineState, covering: \(targetParameter.name).size\(threadgroupExpressionString)) }") } else { print("Could not generate dispatching over parameter \(argument)") }