Optimize span sort dispatch

6 files changed, 92 insertions, 141 deletions

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: &params, length: MemoryLayout<Params>.stride, options: .storageModeShared)!
+        let paramBuffer = device.makeBuffer(
+            bytes: &params, length: MemoryLayout<Params>.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 <metal_stdlib>
 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<float, access::read>  colorTex [[texture(0)]],
@@ -95,7 +78,7 @@ kernel void createMask(
     maskTex.write(m, gid);
 }
 
-// ---- clear span buffer ---------------------------------------------
+// --- clear span buffer ---
 
 kernel void clearSpanBuffer(
     texture2d<uint, access::write> 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<uint, access::read>  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<float, access::read>  colorTex [[texture(0)]],
-    texture2d<half, access::write>  valTex   [[texture(1)]],
-    constant Params &params                  [[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<float, access::read>  colorTex  [[texture(0)]],
-    texture2d<half, access::read>   valTex    [[texture(1)]],
     texture2d<uint, access::read>   spanTex   [[texture(2)]],
     texture2d<float, access::write> sortedTex [[texture(3)]],
     constant Params &params                   [[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;
+        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; }
+        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<uint, access::read>  maskTex     [[texture(0)]],
diff --git a/input.png b/examples/input.png
index 9c9342c..9c9342c 100644
--- a/input.png
+++ b/examples/input.png
diff --git a/examples/output.png b/examples/output.png
new file mode 100644
index 0000000..a29fba9
--- /dev/null
+++ b/examples/output.png
diff --git a/output.png b/output.png
index a29fba9..2be102e 100644
--- a/output.png
+++ b/output.png