Add benchmark scriptHEADmain

3 files changed, 112 insertions, 17 deletions

diff --git a/.gitignore b/.gitignore
index 30bcfa4..5ec3882 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
-.build/
+  .build/
+  .DS_Store
diff --git a/Sources/PixelSort.swift b/Sources/PixelSort.swift
index 139bcca..3980a6d 100644
--- a/Sources/PixelSort.swift
+++ b/Sources/PixelSort.swift
@@ -162,25 +162,67 @@ struct PixelSort: ParsableCommand {
         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)
-        try compute.run(pipeline: createMaskPipeline, width: width, height: height)
 
         identifySpansPipeline.arguments.maskTex = .texture(maskTex)
         identifySpansPipeline.arguments.params = .buffer(paramBuffer)
         identifySpansPipeline.arguments.spanBuffer = .buffer(spanBuffer)
         identifySpansPipeline.arguments.spanCount = .buffer(counterBuffer)
-        try compute.run(pipeline: identifySpansPipeline, width: 1, height: height)
 
         prepareIndirectArgsPipeline.arguments.spanCount = .buffer(counterBuffer)
         prepareIndirectArgsPipeline.arguments.indirectArgs = .buffer(indirectArgsBuffer)
-        try compute.run(pipeline: prepareIndirectArgsPipeline, width: 1, height: 1)
 
-        // Sort each span into sortedTex (one threadgroup per span, bitonic sort)
-        try compute.task(label: "pixelSort") { task in
+        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)
+
+        // Put all GPU work into a single command buffer for timing
+        let gpuStart = CACurrentMediaTime()
+
+        try compute.task(label: "pixelSortFull") { task in
+            // Pass 1: build sort keys
+            try task.run { dispatch in
+                let maxTPT = buildSortKeysPipeline.computePipelineState
+                    .maxTotalThreadsPerThreadgroup
+                let tpgW = Int(sqrt(Double(maxTPT)))
+                let tpgH = maxTPT / tpgW
+                try dispatch(
+                    pipeline: buildSortKeysPipeline,
+                    threads: MTLSize(width: width, height: height, depth: 1),
+                    threadsPerThreadgroup: MTLSize(width: tpgW, height: tpgH, depth: 1))
+            }
+            // Pass 2: create mask
+            try task.run { dispatch in
+                let maxTPT = createMaskPipeline.computePipelineState.maxTotalThreadsPerThreadgroup
+                let tpgW = Int(sqrt(Double(maxTPT)))
+                let tpgH = maxTPT / tpgW
+                try dispatch(
+                    pipeline: createMaskPipeline,
+                    threads: MTLSize(width: width, height: height, depth: 1),
+                    threadsPerThreadgroup: MTLSize(width: tpgW, height: tpgH, depth: 1))
+            }
+            // Pass 3: identify spans
+            try task.run { dispatch in
+                let maxTPT = identifySpansPipeline.computePipelineState
+                    .maxTotalThreadsPerThreadgroup
+                try dispatch(
+                    pipeline: identifySpansPipeline,
+                    threads: MTLSize(width: 1, height: height, depth: 1),
+                    threadsPerThreadgroup: MTLSize(width: min(maxTPT, height), height: 1, depth: 1))
+            }
+            // Pass 4: prepare indirect args
+            try task.run { dispatch in
+                try dispatch(
+                    pipeline: prepareIndirectArgsPipeline,
+                    threads: MTLSize(width: 1, height: 1, depth: 1),
+                    threadsPerThreadgroup: MTLSize(width: 1, height: 1, depth: 1))
+            }
+            // Pass 5: bitonic sort (indirect dispatch)
             try task.run { dispatch in
                 let enc = dispatch.commandEncoder
                 enc.setComputePipelineState(pixelSortPipeline.computePipelineState)
@@ -192,18 +234,21 @@ struct PixelSort: ParsableCommand {
                 enc.dispatchThreadgroups(
                     indirectBuffer: indirectArgsBuffer,
                     indirectBufferOffset: 0,
-                    threadsPerThreadgroup: MTLSize(width: 1024, height: 1, depth: 1)
-                )
+                    threadsPerThreadgroup: MTLSize(width: 1024, height: 1, depth: 1))
+            }
+            // Pass 6: composite
+            try task.run { dispatch in
+                let maxTPT = compositePipeline.computePipelineState.maxTotalThreadsPerThreadgroup
+                let tpgW = Int(sqrt(Double(maxTPT)))
+                let tpgH = maxTPT / tpgW
+                try dispatch(
+                    pipeline: compositePipeline,
+                    threads: MTLSize(width: width, height: height, depth: 1),
+                    threadsPerThreadgroup: MTLSize(width: tpgW, height: tpgH, depth: 1))
             }
         }
-
-        // 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)
+        let gpuMs = (CACurrentMediaTime() - gpuStart) * 1000.0
+        print(String(format: "GPU time: %.3f ms (encode + execute, %d×%d)", gpuMs, width, height))
 
         var outputData = [UInt8](repeating: 0, count: width * height * bytesPerPixel)
         texB.getBytes(
diff --git a/bench.sh b/bench.sh
new file mode 100755
index 0000000..049b035
--- /dev/null
+++ b/bench.sh
@@ -0,0 +1,49 @@
+#!/bin/bash
+set -euo pipefail
+
+BIN=".build/release/pixel-sort"
+INPUT="examples/input-1920x1080.png"
+OUTPUT="/tmp/pixel-sort-bench-output.png"
+ITERS="${1:-100}"
+
+echo "Binary:     $BIN"
+echo "Input:      $INPUT"
+echo "Iterations: $ITERS"
+echo ""
+
+# Warmup
+"$BIN" "$INPUT" "$OUTPUT" >/dev/null 2>&1
+
+times=()
+for i in $(seq 1 "$ITERS"); do
+    t=$( { /usr/bin/time -p "$BIN" "$INPUT" "$OUTPUT" ; } 2>&1 | awk '/^real/ {print $2}')
+    times+=("$t")
+done
+
+# Sort and compute stats
+sorted=($(printf '%s\n' "${times[@]}" | sort -n))
+n=${#sorted[@]}
+min=${sorted[0]}
+max=${sorted[$((n-1))]}
+median=${sorted[$((n/2))]}
+p95=${sorted[$(( (n * 95) / 100 ))]}
+p99=${sorted[$(( (n * 99) / 100 ))]}
+
+sum=$(printf '%s\n' "${times[@]}" | awk '{s+=$1} END {print s}')
+mean=$(echo "$sum $n" | awk '{printf "%.4f", $1/$2}')
+
+echo "=== Results (1920x1080, ${ITERS} runs) ==="
+echo "  Min:    ${min}s"
+echo "  Mean:   ${mean}s"
+echo "  Median: ${median}s"
+echo "  P95:    ${p95}s"
+echo "  P99:    ${p99}s"
+echo "  Max:    ${max}s"
+echo ""
+echo "  Target: ≤0.002s (2ms) P95"
+
+p95_ms=$(echo "$p95" | awk '{printf "%.1f", $1 * 1000}')
+pass=$(echo "$p95" | awk '{print ($1 <= 0.002) ? "PASS" : "FAIL"}')
+echo "  P95 = ${p95_ms}ms → ${pass}"
+
+rm -f "$OUTPUT"