Add benchmark scriptHEADmain

1 files changed, 61 insertions, 16 deletions

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(