diff options
Diffstat (limited to 'Sources/PixelSort.swift')
| -rw-r--r-- | Sources/PixelSort.swift | 171 |
1 files changed, 107 insertions, 64 deletions
diff --git a/Sources/PixelSort.swift b/Sources/PixelSort.swift index dff52f2..fa1bee6 100644 --- a/Sources/PixelSort.swift +++ b/Sources/PixelSort.swift @@ -24,9 +24,18 @@ struct PixelSort: ParsableCommand { @Option(name: .shortAndLong, help: "Upper brightness threshold (0.0–1.0)") var upper: Float = 0.9 - @Flag(name: .shortAndLong, help: "Sort descending") + @Flag(name: .shortAndLong, help: "Sort descending (reverse)") var descending: Bool = false + @Option(name: .shortAndLong, help: "Gamma applied to sorted pixels (Unity-style composite)") + var gamma: Float = 1.0 + + @Option(name: .shortAndLong, help: "Clamp maximum sortable span length (default: image width)") + var maxSpan: Int? + + @Flag(name: .long, help: "Invert the threshold mask") + var invertMask: Bool = false + mutating func run() throws { let device = MTLCreateSystemDefaultDevice()! let compute = try Compute(device: device) @@ -44,18 +53,51 @@ struct PixelSort: ParsableCommand { let width = cgImage.width let height = cgImage.height - let desc = MTLTextureDescriptor.texture2DDescriptor( + let rgbaDesc = MTLTextureDescriptor.texture2DDescriptor( pixelFormat: .rgba8Unorm, width: width, height: height, mipmapped: false ) - desc.usage = [.shaderRead, .shaderWrite] + rgbaDesc.usage = [.shaderRead, .shaderWrite] + + let texA = device.makeTexture(descriptor: rgbaDesc)! + let texB = device.makeTexture(descriptor: rgbaDesc)! + texA.label = "original" + texB.label = "output" - let texA = device.makeTexture(descriptor: desc)! - let texB = device.makeTexture(descriptor: desc)! - texA.label = "texA" - texB.label = "texB" + let sortedTex = device.makeTexture(descriptor: rgbaDesc)! + sortedTex.label = "sorted" + + let maskDesc = MTLTextureDescriptor.texture2DDescriptor( + pixelFormat: .r8Uint, + width: width, + height: height, + mipmapped: false + ) + maskDesc.usage = [.shaderRead, .shaderWrite] + let maskTex = device.makeTexture(descriptor: maskDesc)! + maskTex.label = "mask" + + let spanDesc = MTLTextureDescriptor.texture2DDescriptor( + pixelFormat: .r32Uint, + width: width, + height: height, + mipmapped: false + ) + spanDesc.usage = [.shaderRead, .shaderWrite] + let spanTex = device.makeTexture(descriptor: spanDesc)! + spanTex.label = "spans" + + let valDesc = MTLTextureDescriptor.texture2DDescriptor( + pixelFormat: .r16Float, + width: width, + height: height, + mipmapped: false + ) + valDesc.usage = [.shaderRead, .shaderWrite] + let valTex = device.makeTexture(descriptor: valDesc)! + valTex.label = "sortValues" // Upload pixel data let bytesPerPixel = 4 @@ -86,58 +128,68 @@ struct PixelSort: ParsableCommand { let shaderSource = try String(contentsOf: Bundle.module.url(forResource: "Shaders", withExtension: "metal")!, encoding: .utf8) let library = ShaderLibrary.source(shaderSource) - var copyPipeline = try compute.makePipeline(function: library.copyTexture) - var sortPipeline = try compute.makePipeline(function: library.bitonicSortStep) + var createMaskPipeline = try compute.makePipeline(function: library.createMask) + var clearSpanPipeline = try compute.makePipeline(function: library.clearSpanBuffer) + var identifySpansPipeline = try compute.makePipeline(function: library.identifySpans) + var rgbToValPipeline = try compute.makePipeline(function: library.rgbToSortValue) + var pixelSortPipeline = try compute.makePipeline(function: library.pixelSortSpan) + var compositePipeline = try compute.makePipeline(function: library.composite) - // Bitonic sort requires log2(nextPow2(width)) outer passes - let n = nextPowerOf2(width) - - // Ping-pong between texA and texB - var readTex = texA - var writeTex = texB - - var blockSize: Int = 2 - while blockSize <= n { - var subBlockSize = blockSize - while subBlockSize >= 2 { - // Copy readTex → writeTex so untouched pixels carry forward - copyPipeline.arguments.src = .texture(readTex) - copyPipeline.arguments.dst = .texture(writeTex) - try compute.run(pipeline: copyPipeline, width: width, height: height) + var params = Params( + width: UInt32(width), + height: UInt32(height), + sortKey: UInt32(key.metalValue), + lowerThreshold: lower, + upperThreshold: upper, + reverseSorting: descending ? 1 : 0, + gamma: gamma, + maxSpanLength: UInt32(maxSpan ?? width), + invertMask: invertMask ? 1 : 0 + ) + let paramBuffer = device.makeBuffer(bytes: ¶ms, length: MemoryLayout<Params>.stride, options: .storageModeShared)! - // Run the bitonic compare-swap step - var params = BitonicParams( - width: UInt32(width), - height: UInt32(height), - blockSize: UInt32(blockSize), - subBlockSize: UInt32(subBlockSize), - sortKey: UInt32(key.metalValue), - lowerThreshold: lower, - upperThreshold: upper, - descending: descending ? 1 : 0 - ) + // 1) Mask + createMaskPipeline.arguments.colorTex = .texture(texA) + createMaskPipeline.arguments.maskTex = .texture(maskTex) + createMaskPipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: createMaskPipeline, width: width, height: height) - sortPipeline.arguments.inputTexture = .texture(readTex) - sortPipeline.arguments.outputTexture = .texture(writeTex) + // 2) Clear span buffer + clearSpanPipeline.arguments.spanTex = .texture(spanTex) + clearSpanPipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: clearSpanPipeline, width: width, height: height) - let paramBuffer = device.makeBuffer(bytes: ¶ms, length: MemoryLayout<BitonicParams>.stride, options: .storageModeShared)! - sortPipeline.arguments.params = .buffer(paramBuffer) + // 3) Identify spans (1 thread per row: dispatch width=1) + identifySpansPipeline.arguments.maskTex = .texture(maskTex) + identifySpansPipeline.arguments.spanTex = .texture(spanTex) + identifySpansPipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: identifySpansPipeline, width: 1, height: height) - try compute.run(pipeline: sortPipeline, width: width / 2, height: height) + // 4) Sort values + rgbToValPipeline.arguments.colorTex = .texture(texA) + rgbToValPipeline.arguments.valTex = .texture(valTex) + rgbToValPipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: rgbToValPipeline, width: width, height: height) - // Swap - let tmp = readTex - readTex = writeTex - writeTex = tmp + // 5) Sort each span into sortedTex + pixelSortPipeline.arguments.colorTex = .texture(texA) + pixelSortPipeline.arguments.valTex = .texture(valTex) + pixelSortPipeline.arguments.spanTex = .texture(spanTex) + pixelSortPipeline.arguments.sortedTex = .texture(sortedTex) + pixelSortPipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: pixelSortPipeline, width: width, height: height) - subBlockSize /= 2 - } - blockSize *= 2 - } + // 6) Composite only masked pixels into output texB + compositePipeline.arguments.maskTex = .texture(maskTex) + compositePipeline.arguments.sortedTex = .texture(sortedTex) + compositePipeline.arguments.originalTex = .texture(texA) + compositePipeline.arguments.outTex = .texture(texB) + compositePipeline.arguments.params = .buffer(paramBuffer) + try compute.run(pipeline: compositePipeline, width: width, height: height) - // Read back from readTex (the last write destination after swap) + // Read back from output var outputData = [UInt8](repeating: 0, count: width * height * bytesPerPixel) - readTex.getBytes( + texB.getBytes( &outputData, bytesPerRow: bytesPerRow, from: MTLRegionMake2D(0, 0, width, height), @@ -173,25 +225,16 @@ struct PixelSort: ParsableCommand { // MARK: - Helpers -struct BitonicParams { +struct Params { var width: UInt32 var height: UInt32 - var blockSize: UInt32 - var subBlockSize: UInt32 var sortKey: UInt32 var lowerThreshold: Float var upperThreshold: Float - var descending: UInt32 -} - -func nextPowerOf2(_ n: Int) -> Int { - var v = n - 1 - v |= v >> 1 - v |= v >> 2 - v |= v >> 4 - v |= v >> 8 - v |= v >> 16 - return v + 1 + var reverseSorting: UInt32 + var gamma: Float + var maxSpanLength: UInt32 + var invertMask: UInt32 } enum SortKeyOption: String, ExpressibleByArgument, CaseIterable { |