From ef7ee12c4af07387a4961fd5c2e0314712d408c1 Mon Sep 17 00:00:00 2001 From: y-jan137 Date: Wed, 29 Apr 2026 08:35:11 +0300 Subject: Add precompute sort key tex before sort --- Sources/PixelSort.swift | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) (limited to 'Sources/PixelSort.swift') 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.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( -- cgit v1.2.3