Add precompute sort key tex before sort

2 files changed, 36 insertions, 5 deletions

diff --git a/Sources/PixelSort.swift b/Sources/PixelSort.swift
index 066baf1..7fc1ce2 100644
--- a/Sources/PixelSort.swift
+++ b/Sources/PixelSort.swift
@@ -78,6 +78,16 @@ struct PixelSort: ParsableCommand {
         let maskTex = device.makeTexture(descriptor: maskDesc)!
         maskTex.label = "mask"
 
+        let sortKeyDesc = MTLTextureDescriptor.texture2DDescriptor(
+            pixelFormat: .r32Float,
+            width: width,
+            height: height,
+            mipmapped: false
+        )
+        sortKeyDesc.usage = [.shaderRead, .shaderWrite]
+        let sortKeyTex = device.makeTexture(descriptor: sortKeyDesc)!
+        sortKeyTex.label = "sortKeys"
+
         // Span descriptor buffer (max one span per pixel is a safe upper bound)
         let maxSpans = width * height
         let spanBufferSize = maxSpans * MemoryLayout<SpanDescriptor>.stride
@@ -126,6 +136,7 @@ struct PixelSort: ParsableCommand {
         let library = ShaderLibrary.source(shaderSource)
 
         var createMaskPipeline = try compute.makePipeline(function: library.createMask)
+        var buildSortKeysPipeline = try compute.makePipeline(function: library.buildSortKeys)
         var identifySpansPipeline = try compute.makePipeline(function: library.identifySpans)
         var prepareIndirectArgsPipeline = try compute.makePipeline(
             function: library.prepareIndirectArgs)
@@ -148,6 +159,11 @@ struct PixelSort: ParsableCommand {
 
         counterBuffer.contents().assumingMemoryBound(to: UInt32.self).pointee = 0
 
+        buildSortKeysPipeline.arguments.colorTex = .texture(texA)
+        buildSortKeysPipeline.arguments.sortKeyTex = .texture(sortKeyTex)
+        buildSortKeysPipeline.arguments.params = .buffer(paramBuffer)
+        try compute.run(pipeline: buildSortKeysPipeline, width: width, height: height)
+
         createMaskPipeline.arguments.colorTex = .texture(texA)
         createMaskPipeline.arguments.maskTex = .texture(maskTex)
         createMaskPipeline.arguments.params = .buffer(paramBuffer)
@@ -169,7 +185,8 @@ struct PixelSort: ParsableCommand {
                 let enc = dispatch.commandEncoder
                 enc.setComputePipelineState(pixelSortPipeline.computePipelineState)
                 enc.setTexture(texA, index: 0)
-                enc.setTexture(sortedTex, index: 1)
+                enc.setTexture(sortKeyTex, index: 1)
+                enc.setTexture(sortedTex, index: 2)
                 enc.setBuffer(paramBuffer, offset: 0, index: 0)
                 enc.setBuffer(spanBuffer, offset: 0, index: 1)
                 enc.dispatchThreadgroups(
diff --git a/Sources/Shaders.metal b/Sources/Shaders.metal
index 832ccff..c37cc40 100644
--- a/Sources/Shaders.metal
+++ b/Sources/Shaders.metal
@@ -78,6 +78,20 @@ kernel void createMask(
     maskTex.write(m, gid);
 }
 
+// --- full-frame sort-by scalars (one float per pixel) ---
+
+kernel void buildSortKeys(
+    texture2d<float, access::read>   colorTex   [[texture(0)]],
+    texture2d<float, access::write> sortKeyTex  [[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);
+    sortKeyTex.write(float4(v, 0.0f, 0.0f, 0.0f), gid);
+}
+
 // --- span descriptor for indirect dispatch ---
 
 struct SpanDescriptor {
@@ -151,8 +165,9 @@ kernel void prepareIndirectArgs(
 constant uint MAX_LOCAL_SPAN = 2048;
 
 kernel void pixelSortSpan(
-    texture2d<float, access::read>  colorTex    [[texture(0)]],
-    texture2d<float, access::write> sortedTex   [[texture(1)]],
+    texture2d<float, access::read>  colorTex     [[texture(0)]],
+    texture2d<float, access::read> sortKeyTex   [[texture(1)]],
+    texture2d<float, access::write> sortedTex   [[texture(2)]],
     constant Params &params                     [[buffer(0)]],
     device const SpanDescriptor *spanBuffer     [[buffer(1)]],
     uint gid                                    [[thread_position_in_grid]]
@@ -166,8 +181,7 @@ kernel void pixelSortSpan(
 
     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);
+        cache[k] = sortKeyTex.read(uint2(x + k, row)).x;
     }
 
     float minValue = cache[0];