aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--README.md8
-rw-r--r--Sources/PixelSort.swift75
-rw-r--r--Sources/Shaders.metal138
-rw-r--r--examples/input-1920x1080.jpgbin0 -> 294494 bytes
-rw-r--r--examples/input-626x417.png (renamed from examples/input.png)bin45495 -> 45495 bytes
-rw-r--r--examples/output-1920-1080.jpgbin0 -> 438790 bytes
-rw-r--r--examples/output-626x417.png (renamed from examples/output.png)bin424216 -> 424216 bytes
-rw-r--r--output.pngbin422068 -> 0 bytes
8 files changed, 133 insertions, 88 deletions
diff --git a/README.md b/README.md
index e9f317e..ef7e82e 100644
--- a/README.md
+++ b/README.md
@@ -2,3 +2,11 @@
swift build -c release
./.build/release/pixel-sort input.png output2.png --lower 0.2 --upper 0.8 --key brightness
```
+
+Performance target: 2ms max for 1080p
+
+Shader optimization bottlenecks:
+- VRAM
+- Shader code itself
+ - Texture sampling (memory bandwidth): precalculate sorting value (SortKey)
+ -
diff --git a/Sources/PixelSort.swift b/Sources/PixelSort.swift
index a9ba590..066baf1 100644
--- a/Sources/PixelSort.swift
+++ b/Sources/PixelSort.swift
@@ -78,15 +78,21 @@ struct PixelSort: ParsableCommand {
let maskTex = device.makeTexture(descriptor: maskDesc)!
maskTex.label = "mask"
- let spanDesc = MTLTextureDescriptor.texture2DDescriptor(
- pixelFormat: .r32Uint,
- width: width,
- height: height,
- mipmapped: false
- )
- spanDesc.usage = [.shaderRead, .shaderWrite]
- let spanTex = device.makeTexture(descriptor: spanDesc)!
- spanTex.label = "spans"
+ // Span descriptor buffer (max one span per pixel is a safe upper bound)
+ let maxSpans = width * height
+ let spanBufferSize = maxSpans * MemoryLayout<SpanDescriptor>.stride
+ let spanBuffer = device.makeBuffer(length: spanBufferSize, options: .storageModeShared)!
+ spanBuffer.label = "spanBuffer"
+
+ // Atomic counter for number of spans found
+ let counterBuffer = device.makeBuffer(
+ length: MemoryLayout<UInt32>.stride, options: .storageModeShared)!
+ counterBuffer.label = "spanCount"
+
+ // Indirect dispatch arguments buffer (3 x uint32)
+ let indirectArgsBuffer = device.makeBuffer(
+ length: MemoryLayout<UInt32>.stride * 3, options: .storageModeShared)!
+ indirectArgsBuffer.label = "indirectArgs"
let bytesPerPixel = 4
let bytesPerRow = bytesPerPixel * width
@@ -120,9 +126,10 @@ struct PixelSort: ParsableCommand {
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 pixelSortPipeline = try compute.makePipeline(function: library.pixelSortSpan)
+ var prepareIndirectArgsPipeline = try compute.makePipeline(
+ function: library.prepareIndirectArgs)
+ let pixelSortPipeline = try compute.makePipeline(function: library.pixelSortSpan)
var compositePipeline = try compute.makePipeline(function: library.composite)
var params = Params(
@@ -139,31 +146,41 @@ struct PixelSort: ParsableCommand {
let paramBuffer = device.makeBuffer(
bytes: &params, length: MemoryLayout<Params>.stride, options: .storageModeShared)!
- // Mask
+ counterBuffer.contents().assumingMemoryBound(to: UInt32.self).pointee = 0
+
createMaskPipeline.arguments.colorTex = .texture(texA)
createMaskPipeline.arguments.maskTex = .texture(maskTex)
createMaskPipeline.arguments.params = .buffer(paramBuffer)
try compute.run(pipeline: createMaskPipeline, width: width, height: height)
- // Clear span buffer
- clearSpanPipeline.arguments.spanTex = .texture(spanTex)
- clearSpanPipeline.arguments.params = .buffer(paramBuffer)
- try compute.run(pipeline: clearSpanPipeline, width: width, height: height)
-
- // Identify spans (1 thread per row: dispatch width=1)
identifySpansPipeline.arguments.maskTex = .texture(maskTex)
- identifySpansPipeline.arguments.spanTex = .texture(spanTex)
identifySpansPipeline.arguments.params = .buffer(paramBuffer)
+ identifySpansPipeline.arguments.spanBuffer = .buffer(spanBuffer)
+ identifySpansPipeline.arguments.spanCount = .buffer(counterBuffer)
try compute.run(pipeline: identifySpansPipeline, width: 1, height: height)
- // Sort each span into sortedTex
- pixelSortPipeline.arguments.colorTex = .texture(texA)
- pixelSortPipeline.arguments.spanTex = .texture(spanTex)
- pixelSortPipeline.arguments.sortedTex = .texture(sortedTex)
- pixelSortPipeline.arguments.params = .buffer(paramBuffer)
- try compute.run(pipeline: pixelSortPipeline, 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 thread per span)
+ try compute.task(label: "pixelSort") { task in
+ try task.run { dispatch in
+ let enc = dispatch.commandEncoder
+ enc.setComputePipelineState(pixelSortPipeline.computePipelineState)
+ enc.setTexture(texA, index: 0)
+ enc.setTexture(sortedTex, index: 1)
+ enc.setBuffer(paramBuffer, offset: 0, index: 0)
+ enc.setBuffer(spanBuffer, offset: 0, index: 1)
+ enc.dispatchThreadgroups(
+ indirectBuffer: indirectArgsBuffer,
+ indirectBufferOffset: 0,
+ threadsPerThreadgroup: MTLSize(width: 1, height: 1, depth: 1)
+ )
+ }
+ }
- // 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)
@@ -206,6 +223,12 @@ struct PixelSort: ParsableCommand {
}
}
+struct SpanDescriptor {
+ var row: UInt32
+ var startX: UInt32
+ var length: UInt32
+}
+
struct Params {
var width: UInt32
var height: UInt32
diff --git a/Sources/Shaders.metal b/Sources/Shaders.metal
index c8989e9..832ccff 100644
--- a/Sources/Shaders.metal
+++ b/Sources/Shaders.metal
@@ -56,7 +56,7 @@ struct Params {
float lowerThreshold; // luminance low
float upperThreshold; // luminance high
uint reverseSorting; // 0 = normal, 1 = reverse
- float gamma; // output gamma (Unity applies pow(abs(sorted), gamma))
+ float gamma; // output gamma
uint maxSpanLength; // clamp span length (safety)
uint invertMask; // 0/1
};
@@ -78,24 +78,22 @@ kernel void createMask(
maskTex.write(m, gid);
}
-// --- clear span buffer ---
+// --- span descriptor for indirect dispatch ---
-kernel void clearSpanBuffer(
- texture2d<uint, access::write> spanTex [[texture(0)]],
- constant Params &params [[buffer(0)]],
- uint2 gid [[thread_position_in_grid]]
-) {
- if (gid.x >= params.width || gid.y >= params.height) return;
- spanTex.write(0u, gid);
-}
+struct SpanDescriptor {
+ uint row;
+ uint startX;
+ uint length;
+};
-// --- identify spans w mask ---
+// --- identify spans and build span buffer ---
kernel void identifySpans(
- texture2d<uint, access::read> maskTex [[texture(0)]],
- texture2d<uint, access::write> spanTex [[texture(1)]],
- constant Params &params [[buffer(0)]],
- uint2 gid [[thread_position_in_grid]]
+ texture2d<uint, access::read> maskTex [[texture(0)]],
+ constant Params &params [[buffer(0)]],
+ device SpanDescriptor *spanBuffer [[buffer(1)]],
+ device atomic_uint *spanCount [[buffer(2)]],
+ uint2 gid [[thread_position_in_grid]]
) {
uint row = gid.y;
if (gid.x != 0 || row >= params.height) return;
@@ -112,7 +110,8 @@ kernel void identifySpans(
if (m == 0 || spanLength >= spanLimit) {
if (spanLength != 0) {
uint outLen = (m == 1u) ? (spanLength + 1u) : spanLength;
- spanTex.write(outLen, uint2(spanStart, row));
+ uint idx = atomic_fetch_add_explicit(spanCount, 1, memory_order_relaxed);
+ spanBuffer[idx] = SpanDescriptor { row, spanStart, outLen };
}
spanStart = pos;
spanLength = 0;
@@ -122,68 +121,83 @@ kernel void identifySpans(
}
if (spanLength != 0 && spanStart < params.width) {
- spanTex.write(spanLength, uint2(spanStart, row));
+ uint idx = atomic_fetch_add_explicit(spanCount, 1, memory_order_relaxed);
+ spanBuffer[idx] = SpanDescriptor { row, spanStart, spanLength };
}
}
-// --- sort pixels within a span ---
+// --- prepare indirect dispatch arguments from span count ---
-constant uint MAX_LOCAL_SPAN = 2048;
+struct IndirectArgs {
+ uint threadgroupsX;
+ uint threadgroupsY;
+ uint threadgroupsZ;
+};
-kernel void pixelSortSpan(
- texture2d<float, access::read> colorTex [[texture(0)]],
- 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]]
+kernel void prepareIndirectArgs(
+ device atomic_uint *spanCount [[buffer(0)]],
+ device IndirectArgs *indirectArgs [[buffer(1)]],
+ uint gid [[thread_position_in_grid]]
) {
- uint row = gid.y;
- if (gid.x != 0 || row >= params.height) return;
+ if (gid != 0) return;
+ uint count = atomic_load_explicit(spanCount, memory_order_relaxed);
+ indirectArgs->threadgroupsX = count;
+ indirectArgs->threadgroupsY = 1;
+ indirectArgs->threadgroupsZ = 1;
+}
- for (uint x = 0; x < params.width; ++x) {
- uint spanLength = spanTex.read(uint2(x, row)).x;
- if (spanLength == 0) continue;
+// --- sort pixels within a span (one thread per span, indirect dispatch) ---
- spanLength = min(spanLength, params.width - x);
- spanLength = min(spanLength, max(1u, params.maxSpanLength));
- spanLength = min(spanLength, MAX_LOCAL_SPAN);
+constant uint MAX_LOCAL_SPAN = 2048;
- 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);
- }
+kernel void pixelSortSpan(
+ texture2d<float, access::read> colorTex [[texture(0)]],
+ texture2d<float, access::write> sortedTex [[texture(1)]],
+ constant Params &params [[buffer(0)]],
+ device const SpanDescriptor *spanBuffer [[buffer(1)]],
+ uint gid [[thread_position_in_grid]]
+) {
+ SpanDescriptor span = spanBuffer[gid];
+ uint row = span.row;
+ uint x = span.startX;
+ uint spanLength = min(span.length, params.width - x);
+ spanLength = min(spanLength, max(1u, params.maxSpanLength));
+ spanLength = min(spanLength, MAX_LOCAL_SPAN);
+
+ 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];
- 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(x + minIndex, row));
- float4 cMax = colorTex.read(uint2(x + maxIndex, row));
+ float4 cMin = colorTex.read(uint2(x + minIndex, row));
+ float4 cMax = colorTex.read(uint2(x + maxIndex, row));
- sortedTex.write(cMin, uint2(x + dstMin, row));
- sortedTex.write(cMax, uint2(x + dstMax, row));
+ 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;
- }
- x += (spanLength - 1);
+ cache[minIndex] = 2.0f;
+ cache[maxIndex] = -2.0f;
+ minValue = 1.0f;
+ maxValue = -1.0f;
}
}
diff --git a/examples/input-1920x1080.jpg b/examples/input-1920x1080.jpg
new file mode 100644
index 0000000..98091b0
--- /dev/null
+++ b/examples/input-1920x1080.jpg
Binary files differ
diff --git a/examples/input.png b/examples/input-626x417.png
index 9c9342c..9c9342c 100644
--- a/examples/input.png
+++ b/examples/input-626x417.png
Binary files differ
diff --git a/examples/output-1920-1080.jpg b/examples/output-1920-1080.jpg
new file mode 100644
index 0000000..c0b4f07
--- /dev/null
+++ b/examples/output-1920-1080.jpg
Binary files differ
diff --git a/examples/output.png b/examples/output-626x417.png
index a29fba9..a29fba9 100644
--- a/examples/output.png
+++ b/examples/output-626x417.png
Binary files differ
diff --git a/output.png b/output.png
deleted file mode 100644
index 2be102e..0000000
--- a/output.png
+++ /dev/null
Binary files differ
generated by cgit v1.2.3 (git 2.51.0) at 2026-09-01 19:58:32 +0000