From 5c2531da748443bf10f1bbed90081e6c59a1fd6e Mon Sep 17 00:00:00 2001 From: y-jan137 Date: Mon, 27 Apr 2026 18:47:16 +0300 Subject: Optimize span sort dispatch --- README.md | 4 ++ Sources/PixelSort.swift | 105 +++++++++++++++++----------------------- Sources/Shaders.metal | 126 ++++++++++++++++++------------------------------ examples/input.png | Bin 0 -> 45495 bytes examples/output.png | Bin 0 -> 424216 bytes input.png | Bin 45495 -> 0 bytes output.png | Bin 424216 -> 422068 bytes 7 files changed, 93 insertions(+), 142 deletions(-) create mode 100644 README.md create mode 100644 examples/input.png create mode 100644 examples/output.png delete mode 100644 input.png diff --git a/README.md b/README.md new file mode 100644 index 0000000..e9f317e --- /dev/null +++ b/README.md @@ -0,0 +1,4 @@ +```bash +swift build -c release +./.build/release/pixel-sort input.png output2.png --lower 0.2 --upper 0.8 --key brightness +``` diff --git a/Sources/PixelSort.swift b/Sources/PixelSort.swift index fa1bee6..a9ba590 100644 --- a/Sources/PixelSort.swift +++ b/Sources/PixelSort.swift @@ -6,7 +6,7 @@ import Metal @main struct PixelSort: ParsableCommand { static let configuration = CommandConfiguration( - abstract: "GPU-accelerated pixel sorting for glitch art" + abstract: "Metal-accelerated pixel sorting library" ) @Argument(help: "Input image path") @@ -15,7 +15,7 @@ struct PixelSort: ParsableCommand { @Argument(help: "Output image path") var output: String - @Option(name: .shortAndLong, help: "Sort key: brightness, hue, saturation, red, green, blue") + @Option(name: .shortAndLong, help: "Sort key: brightness, hue, saturation, R, G, B") var key: SortKeyOption = .brightness @Option(name: .shortAndLong, help: "Lower brightness threshold (0.0–1.0)") @@ -40,14 +40,13 @@ struct PixelSort: ParsableCommand { let device = MTLCreateSystemDefaultDevice()! let compute = try Compute(device: device) - // Load image into a Metal texture let inputURL = URL(fileURLWithPath: input) let outputURL = URL(fileURLWithPath: self.output) guard let nsImage = NSImage(contentsOf: inputURL), - let cgImage = nsImage.cgImage(forProposedRect: nil, context: nil, hints: nil) + let cgImage = nsImage.cgImage(forProposedRect: nil, context: nil, hints: nil) else { - throw ValidationError("Could not load image at \(input)") + throw ValidationError("Could not load image") } let width = cgImage.width @@ -89,30 +88,21 @@ struct PixelSort: ParsableCommand { 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 let bytesPerRow = bytesPerPixel * width var pixelData = [UInt8](repeating: 0, count: width * height * bytesPerPixel) let colorSpace = CGColorSpaceCreateDeviceRGB() - guard let ctx = CGContext( - data: &pixelData, - width: width, - height: height, - bitsPerComponent: 8, - bytesPerRow: bytesPerRow, - space: colorSpace, - bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue - ) else { + guard + let ctx = CGContext( + data: &pixelData, + width: width, + height: height, + bitsPerComponent: 8, + bytesPerRow: bytesPerRow, + space: colorSpace, + bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue + ) + else { throw ValidationError("Failed to create CGContext") } ctx.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height)) @@ -124,14 +114,14 @@ struct PixelSort: ParsableCommand { bytesPerRow: bytesPerRow ) - // Load shaders - let shaderSource = try String(contentsOf: Bundle.module.url(forResource: "Shaders", withExtension: "metal")!, encoding: .utf8) + let shaderSource = try String( + contentsOf: Bundle.module.url(forResource: "Shaders", withExtension: "metal")!, + encoding: .utf8) let library = ShaderLibrary.source(shaderSource) 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) @@ -146,40 +136,34 @@ struct PixelSort: ParsableCommand { maxSpanLength: UInt32(maxSpan ?? width), invertMask: invertMask ? 1 : 0 ) - let paramBuffer = device.makeBuffer(bytes: ¶ms, length: MemoryLayout.stride, options: .storageModeShared)! + let paramBuffer = device.makeBuffer( + bytes: ¶ms, length: MemoryLayout.stride, options: .storageModeShared)! - // 1) Mask + // 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) - // 2) Clear span buffer + // Clear span buffer clearSpanPipeline.arguments.spanTex = .texture(spanTex) clearSpanPipeline.arguments.params = .buffer(paramBuffer) try compute.run(pipeline: clearSpanPipeline, width: width, height: height) - // 3) Identify spans (1 thread per row: dispatch width=1) + // 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) - // 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) - - // 5) Sort each span into sortedTex + // 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) + try compute.run(pipeline: pixelSortPipeline, width: 1, height: height) - // 6) Composite only masked pixels into output texB + // Composite only masked pixels into output texB compositePipeline.arguments.maskTex = .texture(maskTex) compositePipeline.arguments.sortedTex = .texture(sortedTex) compositePipeline.arguments.originalTex = .texture(texA) @@ -187,7 +171,6 @@ struct PixelSort: ParsableCommand { compositePipeline.arguments.params = .buffer(paramBuffer) try compute.run(pipeline: compositePipeline, width: width, height: height) - // Read back from output var outputData = [UInt8](repeating: 0, count: width * height * bytesPerPixel) texB.getBytes( &outputData, @@ -196,17 +179,17 @@ struct PixelSort: ParsableCommand { mipmapLevel: 0 ) - // Save output - guard let outCtx = CGContext( - data: &outputData, - width: width, - height: height, - bitsPerComponent: 8, - bytesPerRow: bytesPerRow, - space: colorSpace, - bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue - ), - let outCGImage = outCtx.makeImage() + guard + let outCtx = CGContext( + data: &outputData, + width: width, + height: height, + bitsPerComponent: 8, + bytesPerRow: bytesPerRow, + space: colorSpace, + bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue + ), + let outCGImage = outCtx.makeImage() else { throw ValidationError("Failed to create output image") } @@ -223,8 +206,6 @@ struct PixelSort: ParsableCommand { } } -// MARK: - Helpers - struct Params { var width: UInt32 var height: UInt32 @@ -242,12 +223,12 @@ enum SortKeyOption: String, ExpressibleByArgument, CaseIterable { var metalValue: Int { switch self { - case .brightness: return 0 - case .hue: return 1 - case .saturation: return 2 - case .red: return 3 - case .green: return 4 - case .blue: return 5 + case .brightness: return 0 + case .hue: return 1 + case .saturation: return 2 + case .red: return 3 + case .green: return 4 + case .blue: return 5 } } } diff --git a/Sources/Shaders.metal b/Sources/Shaders.metal index 31608a6..c8989e9 100644 --- a/Sources/Shaders.metal +++ b/Sources/Shaders.metal @@ -1,21 +1,6 @@ #include using namespace metal; -// -------------------------------------------------------------------- -// Span-based pixel sorting (Unity Pixel-Sorting port) -// -// Pipeline: -// 1) createMask: mark pixels whose luminance is within thresholds -// 2) clearSpanBuffer -// 3) identifySpans: for each row, write span length at span start pixel -// 4) rgbToSortValue: compute per-pixel sort value (R/G/B/L/S/H) -// 5) pixelSortSpan: for each span start, sort pixels within the span -// 6) composite: apply sorted pixels only where mask==1 -// -// This matches the semantics of the reference Unity compute shader: -// contiguous masked regions are sorted independently. -// -------------------------------------------------------------------- - enum SortKey : uint { Brightness = 0, Hue = 1, @@ -25,10 +10,9 @@ enum SortKey : uint { Blue = 5, }; -// ---- helpers ------------------------------------------------------- - +// --- helpers --- static inline float luminance(float3 rgb) { - return dot(rgb, float3(0.299, 0.587, 0.114)); + return dot(rgb, float3(0.299, 0.587, 0.114)); // RGB -> greyscale } static inline float hue(float3 c) { @@ -64,8 +48,7 @@ static inline float sort_value(float3 rgb, uint key) { return luminance(rgb); } -// ---- parameters passed from the CPU side --------------------------- - +// --- parameters passed from CPU --- struct Params { uint width; uint height; @@ -78,7 +61,7 @@ struct Params { uint invertMask; // 0/1 }; -// ---- create mask (0/1) --------------------------------------------- +// --- create mask (0/1) --- kernel void createMask( texture2d colorTex [[texture(0)]], @@ -95,7 +78,7 @@ kernel void createMask( maskTex.write(m, gid); } -// ---- clear span buffer --------------------------------------------- +// --- clear span buffer --- kernel void clearSpanBuffer( texture2d spanTex [[texture(0)]], @@ -106,8 +89,7 @@ kernel void clearSpanBuffer( spanTex.write(0u, gid); } -// ---- identify spans (horizontal only) ------------------------------ -// One thread per row: write span length at each span start. +// --- identify spans w mask --- kernel void identifySpans( texture2d maskTex [[texture(0)]], @@ -128,8 +110,6 @@ kernel void identifySpans( pos += 1; if (m == 0 || spanLength >= spanLimit) { - // Write at span start. Mirror Unity behavior: if we hit an unmasked pixel, - // spanLength is current count; if we hit limit while still masked, include current pixel. if (spanLength != 0) { uint outLen = (m == 1u) ? (spanLength + 1u) : spanLength; spanTex.write(outLen, uint2(spanStart, row)); @@ -146,82 +126,68 @@ kernel void identifySpans( } } -// ---- per-pixel sort value ------------------------------------------ - -kernel void rgbToSortValue( - texture2d colorTex [[texture(0)]], - texture2d valTex [[texture(1)]], - constant Params ¶ms [[buffer(0)]], - uint2 gid [[thread_position_in_grid]] -) { - if (gid.x >= params.width || gid.y >= params.height) return; - float3 rgb = saturate(colorTex.read(gid).rgb); - float v = sort_value(rgb, params.sortKey); - valTex.write(half(v), gid); -} - -// ---- sort pixels within a span ------------------------------------- -// One thread per pixel, but only span starts do work. -// Writes sorted pixels into `sortedTex` for the span region. +// --- sort pixels within a span --- constant uint MAX_LOCAL_SPAN = 2048; kernel void pixelSortSpan( texture2d colorTex [[texture(0)]], - texture2d valTex [[texture(1)]], texture2d spanTex [[texture(2)]], texture2d sortedTex [[texture(3)]], constant Params ¶ms [[buffer(0)]], uint2 gid [[thread_position_in_grid]] ) { - if (gid.x >= params.width || gid.y >= params.height) return; + uint row = gid.y; + if (gid.x != 0 || row >= params.height) return; - uint spanLength = spanTex.read(gid).x; - if (spanLength == 0) return; + for (uint x = 0; x < params.width; ++x) { + uint spanLength = spanTex.read(uint2(x, row)).x; + if (spanLength == 0) continue; - spanLength = min(spanLength, params.width - gid.x); - spanLength = min(spanLength, max(1u, params.maxSpanLength)); - spanLength = min(spanLength, MAX_LOCAL_SPAN); + spanLength = min(spanLength, params.width - x); + spanLength = min(spanLength, max(1u, params.maxSpanLength)); + spanLength = min(spanLength, MAX_LOCAL_SPAN); - // Cache sort values for this span. - float cache[MAX_LOCAL_SPAN]; - for (uint k = 0; k < spanLength; ++k) { - cache[k] = float(valTex.read(uint2(gid.x + k, gid.y)).x); - } + float cache[MAX_LOCAL_SPAN]; + for (uint k = 0; k < spanLength; ++k) { + float3 rgb = saturate(colorTex.read(uint2(x + k, row)).rgb); + cache[k] = sort_value(rgb, params.sortKey); + } - float minValue = cache[0]; - float maxValue = cache[0]; - uint minIndex = 0; - uint maxIndex = 0; - - uint steps = (spanLength / 2) + 1; - for (uint i = 0; i < steps; ++i) { - for (uint j = 1; j < spanLength; ++j) { - float v = cache[j]; - // Unity checks `v == saturate(v)` to ignore sentinels; equivalent is 0..1. - if (v >= 0.0f && v <= 1.0f) { - if (v < minValue) { minValue = v; minIndex = j; } - if (maxValue < v) { maxValue = v; maxIndex = j; } + float minValue = cache[0]; + float maxValue = cache[0]; + uint minIndex = 0; + uint maxIndex = 0; + + uint steps = (spanLength / 2) + 1; + for (uint i = 0; i < steps; ++i) { + for (uint j = 1; j < spanLength; ++j) { + float v = cache[j]; + if (v >= 0.0f && v <= 1.0f) { + if (v < minValue) { minValue = v; minIndex = j; } + if (maxValue < v) { maxValue = v; maxIndex = j; } + } } - } - uint dstMin = params.reverseSorting ? i : (spanLength - i - 1); - uint dstMax = params.reverseSorting ? (spanLength - i - 1) : i; + uint dstMin = params.reverseSorting ? i : (spanLength - i - 1); + uint dstMax = params.reverseSorting ? (spanLength - i - 1) : i; - float4 cMin = colorTex.read(uint2(gid.x + minIndex, gid.y)); - float4 cMax = colorTex.read(uint2(gid.x + maxIndex, gid.y)); + float4 cMin = colorTex.read(uint2(x + minIndex, row)); + float4 cMax = colorTex.read(uint2(x + maxIndex, row)); - sortedTex.write(cMin, uint2(gid.x + dstMin, gid.y)); - sortedTex.write(cMax, uint2(gid.x + dstMax, gid.y)); + sortedTex.write(cMin, uint2(x + dstMin, row)); + sortedTex.write(cMax, uint2(x + dstMax, row)); - cache[minIndex] = 2.0f; - cache[maxIndex] = -2.0f; - minValue = 1.0f; - maxValue = -1.0f; + cache[minIndex] = 2.0f; + cache[maxIndex] = -2.0f; + minValue = 1.0f; + maxValue = -1.0f; + } + x += (spanLength - 1); } } -// ---- composite sorted pixels onto original ------------------------- +// --- composite sorted pixels onto original --- kernel void composite( texture2d maskTex [[texture(0)]], diff --git a/examples/input.png b/examples/input.png new file mode 100644 index 0000000..9c9342c Binary files /dev/null and b/examples/input.png differ diff --git a/examples/output.png b/examples/output.png new file mode 100644 index 0000000..a29fba9 Binary files /dev/null and b/examples/output.png differ diff --git a/input.png b/input.png deleted file mode 100644 index 9c9342c..0000000 Binary files a/input.png and /dev/null differ diff --git a/output.png b/output.png index a29fba9..2be102e 100644 Binary files a/output.png and b/output.png differ -- cgit v1.2.3