From 8d6c7e35a09d8a8e351eb9351a55cee612ead80c Mon Sep 17 00:00:00 2001 From: James Jackson-South Date: Thu, 3 Sep 2026 04:31:45 +1000 Subject: [PATCH] Add SIMD AV1 palette clustering --- HEIF_IMPLEMENTATION_PLAN.md | 1 + .../Heif/Av1/Pipeline/Av1PaletteKMeans.cs | 243 ++++++++++++++++++ .../Formats/Heif/Av1/Av1PaletteKMeansTests.cs | 113 ++++++++ 3 files changed, 357 insertions(+) create mode 100644 src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1PaletteKMeans.cs create mode 100644 tests/ImageSharp.Tests/Formats/Heif/Av1/Av1PaletteKMeansTests.cs diff --git a/HEIF_IMPLEMENTATION_PLAN.md b/HEIF_IMPLEMENTATION_PLAN.md index 0b9f3df5c7..96ba858e54 100644 --- a/HEIF_IMPLEMENTATION_PLAN.md +++ b/HEIF_IMPLEMENTATION_PLAN.md @@ -854,6 +854,7 @@ Encoder verification contract: - [~] Luma and chroma palette-color coding now matches current libaom's neighbor-cache flags, sorted delta representation, wrapped V-plane deltas, strict delta-versus-raw V selection, and fixed-point color-rate model at 8, 10, and 12 bits. Encoder costing and emission use only fixed stack spans, including explicitly initialized cache-membership state, and steady-state color costing allocates zero managed bytes. The decoder consumes the same bounded color-syntax primitive after the tile reader derives its neighbor cache, removing duplicated color parsing without changing retained palette ownership. Nine focused syntax, exact palette decode, constrained-allocation, truncation, presentation, and allocation cases pass; all 1,933 entropy cases and all 8,983 HEIF/AV1 cases pass direct net11 Release VSTest. The exact Release rebuild remains at 1,005 warnings and zero errors. Retained encoder palette colors, neighbor caches, color-index maps, candidate generation, and production palette selection remain incomplete, and the compact 8-byte frame mode entries were not enlarged. - [~] Palette color-index map coding now shares the exact current-libaom neighbor weights, stable color ordering, five context classes, first-index uniform code, and diagonal wavefront between encoder costing, encoder writing, and decoder parsing. The decoder's stack-allocated context scores are explicitly cleared before accumulation, removing an invalid dependency on uninitialized stack contents. Costing and writing use a closed generic operation while the shared driver owns traversal and context derivation, so the semantic operations remain independent of map layout and tail handling. The path adds no retained state or per-call managed allocation. Twelve focused map, exact palette decode, padding, trailing-bit, and allocation cases pass; all 1,941 entropy cases and all 8,991 HEIF/AV1 cases pass direct net11 Release VSTest. The exact Release rebuild remains at 1,005 warnings and zero errors. Production payloads remain unchanged because palette selection is still disabled; retained colors, neighbor caches, index-map storage, candidate generation, and production palette mode decision remain incomplete. - [~] Retained encoder palette state and production palette writing now mirror current libaom's 50-byte palette-mode contents, separate luma and shared-chroma sizes, three eight-color planes, above-and-left sorted cache, 64-sample above-cache boundary, mode contexts, palette colors, color-index maps, and syntax order. The current block keeps one inline value in the reusable superblock workspace; only the 4x4-granularity top and left picture edges retain copies for later blocks. For a 3840x2160 tile these edges occupy about 73.4 KiB instead of about 6.2 MiB for a 50-byte palette value attached to every 8x8 mode allocation. Luma and chroma index maps share one lazily allocated 32 KiB owner containing two 128x128 maps, so the palette-disabled production path retains no map owner. The compact final-block workspace falls from about 10.3 KiB to about 8.3 KiB. The writer caps map traversal to the coded plane count, writes maps before transform syntax, and publishes palette edges only after the current block has consumed preceding contexts. Eight focused size, alignment, ownership, cache-boundary, round-trip, map-consumption, and edge-publication cases pass; all 114 palette cases, all 1,942 entropy cases, and all 8,996 HEIF/AV1 cases pass direct foreground net11 Release VSTest. The exact net11 Release rebuild reports 1,050 solution warnings and zero errors; Roslynk reports zero compiler errors and no diagnostics in the touched files. The current-main reference is `a40ed1ea9e4ecc3df58a5bccb76623f2c94ae727`. Production payloads remain unchanged because palette candidate generation is still disabled; that live rate-distortion search is the next checkpoint. +- [~] Luma palette clustering now follows current libaom's one-dimensional search primitive exactly: equal-interval midpoint initialization, first-color tie order, rounded centroid means, deterministic empty-cluster replacement, the 50-iteration limit, and retention of the preceding state when distortion increases. Nearest-color assignment improves on libaom's AVX2 implementation by dispatching Vector512, Vector256, Vector128, then scalar through ImageSharp's shared vector-count helpers. The primitive uses only bounded stack scratch and introduces no allocator rent, managed array, or per-row copy. Three independent tests cover exact centroid convergence, initialization order, 12-bit nearest-color distortion, destination bounds, and every hardware-intrinsic tier. The complete AVIF set passes 8,930 of 8,930 cases and the HEIF set passes 230 of 230 cases through direct foreground net11 Release VSTest. The exact net11 Release rebuild reports 1,050 solution warnings and zero errors, and Roslynk reports zero compiler errors. Candidate enumeration, palette-cache snapping, transform RD selection, and production activation remain in the open luma-palette checkpoint. - [x] The expanded checkpoint exposed a pre-existing transform-block test that asserted uninitialized pooled padding was zero. The test now initializes the complete physical luma plane with a sentinel and proves the block operation leaves both adjacent padding samples unchanged. The exact net11 Release rebuild remains at 1,005 baseline warnings and zero errors, the focused allocator-order set passes 30 of 30 cases, and the complete HEIF/AV1 namespace passes 8,859 of 8,859 direct VSTest cases with zero failures or skips. - [x] Combined-frame OBU output now counts the byte-aligned frame and tile-group headers, non-final tile-size fields, and owned tile payloads before emitting the OBU size. It retains only the small allocator-owned header scratch and writes each entropy-coded tile span directly from its detached owner, removing the second file-sized allocator rent and complete-payload copy. A 64 KiB regression proves exactly one sub-payload-sized byte rent with a balanced return and verifies the exact streamed tile tail; the existing two-tile round trip proves size-prefix and ordering parity. The focused writer and production-frame set passes 32 of 32 direct net11 VSTest cases, current-main `aomdec` accepts all 29 generated native-format payloads, and the complete HEIF/AV1 namespace passes 8,860 of 8,860 cases with zero failures or skips. - [x] Finalized fixed-block decisions now set the block-level transform-skip flag only when every retained luma and coded chroma transform has zero EOB, matching current libaom's conjunction of per-plane skip state. The previous always-false flag produced legal but redundant non-skip and zero-coefficient syntax. Monochrome and 4:2:0 regressions prove both branches from actual coefficient state; the focused decision and production-frame set passes 32 of 32 direct net11 VSTest cases. Current-main `aomdec` accepts all 29 regenerated payloads, the recorded decoded-frame MD5s are unchanged, and affected 16x16 constant 8-bit and 10-bit payloads are one byte smaller. The complete HEIF/AV1 namespace passes 8,862 of 8,862 cases with zero failures or skips. diff --git a/src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1PaletteKMeans.cs b/src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1PaletteKMeans.cs new file mode 100644 index 0000000000..10bf95c26f --- /dev/null +++ b/src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1PaletteKMeans.cs @@ -0,0 +1,243 @@ +// Copyright (c) Six Labors. +// Licensed under the Six Labors Split License. + +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; +using System.Runtime.Intrinsics; + +namespace SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline; + +/// +/// Assigns luma samples to AV1 palette colors and refines one-dimensional palette centroids. +/// +internal static class Av1PaletteKMeans +{ + /// + /// The iteration limit used by the reference encoder for palette clustering. + /// + public const int MaximumIterations = 50; + + /// + /// Assigns every sample to its nearest palette color. + /// + /// The active block samples. + /// The candidate palette colors. + /// The destination palette indices. + /// The sum of squared sample-to-centroid distances. + public static long AssignIndices( + ReadOnlySpan samples, + ReadOnlySpan centroids, + Span indices) + { + Span distanceScratch = stackalloc short[Vector512.Count]; + Span indexScratch = stackalloc short[Vector512.Count]; + ref short sampleBase = ref MemoryMarshal.GetReference(samples); + int offset = 0; + long distortion = 0; + + if (Vector512.IsHardwareAccelerated) + { + nuint vectorCount = Numerics.Vector512Count(samples); + for (; vectorCount > 0; vectorCount--, offset += Vector512.Count) + { + Vector512 sample = Vector512.LoadUnsafe(ref sampleBase, (nuint)offset); + Vector512 bestDistance = Vector512.Abs(sample - Vector512.Create(centroids[0])); + Vector512 bestIndex = Vector512.Zero; + for (int centroidIndex = 1; centroidIndex < centroids.Length; centroidIndex++) + { + Vector512 distance = Vector512.Abs(sample - Vector512.Create(centroids[centroidIndex])); + Vector512 replace = Vector512.LessThan(distance, bestDistance); + bestDistance = Vector512.ConditionalSelect(replace, distance, bestDistance); + bestIndex = Vector512.ConditionalSelect(replace, Vector512.Create((short)centroidIndex), bestIndex); + } + + // Strict comparison preserves the first centroid on ties, matching scalar AV1 palette selection. + bestDistance.CopyTo(distanceScratch); + bestIndex.CopyTo(indexScratch); + for (int lane = 0; lane < Vector512.Count; lane++) + { + indices[offset + lane] = (byte)indexScratch[lane]; + long distance = distanceScratch[lane]; + distortion += distance * distance; + } + } + } + + if (Vector256.IsHardwareAccelerated) + { + nuint vectorCount = Numerics.Vector256Count(samples[offset..]); + for (; vectorCount > 0; vectorCount--, offset += Vector256.Count) + { + Vector256 sample = Vector256.LoadUnsafe(ref sampleBase, (nuint)offset); + Vector256 bestDistance = Vector256.Abs(sample - Vector256.Create(centroids[0])); + Vector256 bestIndex = Vector256.Zero; + for (int centroidIndex = 1; centroidIndex < centroids.Length; centroidIndex++) + { + Vector256 distance = Vector256.Abs(sample - Vector256.Create(centroids[centroidIndex])); + Vector256 replace = Vector256.LessThan(distance, bestDistance); + bestDistance = Vector256.ConditionalSelect(replace, distance, bestDistance); + bestIndex = Vector256.ConditionalSelect(replace, Vector256.Create((short)centroidIndex), bestIndex); + } + + bestDistance.CopyTo(distanceScratch); + bestIndex.CopyTo(indexScratch); + for (int lane = 0; lane < Vector256.Count; lane++) + { + indices[offset + lane] = (byte)indexScratch[lane]; + long distance = distanceScratch[lane]; + distortion += distance * distance; + } + } + } + + if (Vector128.IsHardwareAccelerated) + { + nuint vectorCount = Numerics.Vector128Count(samples[offset..]); + for (; vectorCount > 0; vectorCount--, offset += Vector128.Count) + { + Vector128 sample = Vector128.LoadUnsafe(ref sampleBase, (nuint)offset); + Vector128 bestDistance = Vector128.Abs(sample - Vector128.Create(centroids[0])); + Vector128 bestIndex = Vector128.Zero; + for (int centroidIndex = 1; centroidIndex < centroids.Length; centroidIndex++) + { + Vector128 distance = Vector128.Abs(sample - Vector128.Create(centroids[centroidIndex])); + Vector128 replace = Vector128.LessThan(distance, bestDistance); + bestDistance = Vector128.ConditionalSelect(replace, distance, bestDistance); + bestIndex = Vector128.ConditionalSelect(replace, Vector128.Create((short)centroidIndex), bestIndex); + } + + bestDistance.CopyTo(distanceScratch); + bestIndex.CopyTo(indexScratch); + for (int lane = 0; lane < Vector128.Count; lane++) + { + indices[offset + lane] = (byte)indexScratch[lane]; + long distance = distanceScratch[lane]; + distortion += distance * distance; + } + } + } + + for (; offset < samples.Length; offset++) + { + int bestDistance = Math.Abs(samples[offset] - centroids[0]); + int bestIndex = 0; + for (int centroidIndex = 1; centroidIndex < centroids.Length; centroidIndex++) + { + int distance = Math.Abs(samples[offset] - centroids[centroidIndex]); + if (distance < bestDistance) + { + bestDistance = distance; + bestIndex = centroidIndex; + } + } + + indices[offset] = (byte)bestIndex; + distortion += (long)bestDistance * bestDistance; + } + + return distortion; + } + + /// + /// Refines initialized palette colors through the reference encoder's deterministic clustering sequence. + /// + /// The active block samples. + /// The initialized colors, replaced with the best refined colors. + /// The palette indices belonging to the retained colors. + /// The retained sum of squared distances. + public static long Cluster( + ReadOnlySpan samples, + Span centroids, + Span indices) + { + Span alternateCentroids = stackalloc short[Av1Constants.PaletteMaxSize]; + Span alternateIndices = stackalloc byte[samples.Length]; + alternateCentroids = alternateCentroids[..centroids.Length]; + long distortion = AssignIndices(samples, centroids, indices); + bool currentIsAlternate = false; + + for (int iteration = 0; iteration < MaximumIterations; iteration++) + { + ReadOnlySpan currentCentroids = currentIsAlternate ? alternateCentroids : centroids; + ReadOnlySpan currentIndices = currentIsAlternate ? alternateIndices : indices; + Span nextCentroids = currentIsAlternate ? centroids : alternateCentroids; + Span nextIndices = currentIsAlternate ? indices : alternateIndices; + CalculateCentroids(samples, currentIndices, nextCentroids); + if (nextCentroids.SequenceEqual(currentCentroids)) + { + break; + } + + long nextDistortion = AssignIndices(samples, nextCentroids, nextIndices); + if (nextDistortion > distortion) + { + break; + } + + distortion = nextDistortion; + currentIsAlternate = !currentIsAlternate; + } + + if (currentIsAlternate) + { + alternateCentroids.CopyTo(centroids); + alternateIndices.CopyTo(indices); + } + + return distortion; + } + + /// + /// Places initial colors at the midpoint of equal intervals spanning the sample range. + /// + /// The smallest sample value. + /// The largest sample value. + /// The palette colors to initialize. + public static void InitializeCentroids(short minimum, short maximum, Span centroids) + { + int range = maximum - minimum; + for (int index = 0; index < centroids.Length; index++) + { + centroids[index] = (short)(minimum + (((2 * index) + 1) * range / centroids.Length / 2)); + } + } + + /// + /// Recalculates each centroid from its assigned samples. + /// + private static void CalculateCentroids( + ReadOnlySpan samples, + ReadOnlySpan indices, + Span centroids) + { + Span counts = stackalloc int[Av1Constants.PaletteMaxSize]; + Span sums = stackalloc int[Av1Constants.PaletteMaxSize]; + counts = counts[..centroids.Length]; + sums = sums[..centroids.Length]; + counts.Clear(); + sums.Clear(); + for (int index = 0; index < samples.Length; index++) + { + int centroidIndex = indices[index]; + counts[centroidIndex]++; + sums[centroidIndex] += samples[index]; + } + + uint randomState = (uint)samples[0]; + for (int centroidIndex = 0; centroidIndex < centroids.Length; centroidIndex++) + { + int count = counts[centroidIndex]; + if (count == 0) + { + // Empty clusters use the same seeded sequence on every platform so palette choices remain reproducible. + randomState = unchecked((randomState * 1103515245U) + 12345U); + uint random = (randomState / 65536U) % 32768U; + centroids[centroidIndex] = samples[(int)(random % (uint)samples.Length)]; + } + else + { + centroids[centroidIndex] = (short)((sums[centroidIndex] + (count / 2)) / count); + } + } + } +} diff --git a/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1PaletteKMeansTests.cs b/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1PaletteKMeansTests.cs new file mode 100644 index 0000000000..e92c9d8e7c --- /dev/null +++ b/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1PaletteKMeansTests.cs @@ -0,0 +1,113 @@ +// Copyright (c) Six Labors. +// Licensed under the Six Labors Split License. + +using SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline; +using SixLabors.ImageSharp.Tests.TestUtilities; + +namespace SixLabors.ImageSharp.Tests.Formats.Heif.Av1; + +/// +/// Verifies AV1 palette clustering against independent scalar results at every intrinsic tier. +/// +[Trait("Format", "Avif")] +public class Av1PaletteKMeansTests +{ + /// + /// The hardware configurations covering every descending SIMD width and the scalar fallback. + /// + private const HwIntrinsics Configurations = + HwIntrinsics.AllowAll | HwIntrinsics.DisableAVX512F | HwIntrinsics.DisableAVX | HwIntrinsics.DisableHWIntrinsic; + + /// + /// Verifies nearest-color indices, squared distortion, tie order, and destination bounds. + /// + [Fact] + public void AssignIndicesMatchesScalarAtEveryIntrinsicTier() + => FeatureTestRunner.RunWithHwIntrinsicsFeature(ValidateAssignment, Configurations); + + /// + /// Verifies the deterministic centroid sequence on three separated sample groups. + /// + [Fact] + public void ClusterMatchesReferenceFixture() + { + short[] samples = [0, 1, 2, 100, 101, 102, 200, 201, 202]; + short[] centroids = [33, 100, 167]; + byte[] indices = new byte[samples.Length]; + + long distortion = Av1PaletteKMeans.Cluster(samples, centroids, indices); + + Assert.Equal([1, 101, 201], centroids); + Assert.Equal([0, 0, 0, 1, 1, 1, 2, 2, 2], indices); + Assert.Equal(6, distortion); + } + + /// + /// Verifies the integer interval midpoints used to seed palette refinement. + /// + [Fact] + public void InitializeCentroidsMatchesReferenceIntegerOrder() + { + short[] centroids = new short[3]; + + Av1PaletteKMeans.InitializeCentroids(10, 250, centroids); + + Assert.Equal([50, 130, 210], centroids); + } + + /// + /// Compares production assignment with a scalar equation over a length that exercises every available remainder path. + /// + private static void ValidateAssignment() + { + const int SampleCount = 95; + short[] centroids = [0, 512, 1024, 2048, 3072, 4095]; + short[] samples = new short[SampleCount]; + for (int index = 0; index < samples.Length; index++) + { + samples[index] = (short)(((index * 977) + (index * index * 17)) & 4095); + } + + // This sample is equidistant from the first two colors and must retain the lower palette index. + samples[0] = 256; + byte[] expected = new byte[SampleCount]; + long expectedDistortion = AssignReference(samples, centroids, expected); + byte[] actual = Enumerable.Repeat(byte.MaxValue, SampleCount + 7).ToArray(); + + long actualDistortion = Av1PaletteKMeans.AssignIndices(samples, centroids, actual); + + Assert.Equal(expectedDistortion, actualDistortion); + Assert.Equal(expected, actual.AsSpan(..SampleCount).ToArray()); + Assert.All(actual[SampleCount..], value => Assert.Equal(byte.MaxValue, value)); + } + + /// + /// Applies the scalar nearest-color rule independently of the production SIMD traversal. + /// + private static long AssignReference( + ReadOnlySpan samples, + ReadOnlySpan centroids, + Span indices) + { + long distortion = 0; + for (int sampleIndex = 0; sampleIndex < samples.Length; sampleIndex++) + { + int bestDistance = Math.Abs(samples[sampleIndex] - centroids[0]); + int bestIndex = 0; + for (int centroidIndex = 1; centroidIndex < centroids.Length; centroidIndex++) + { + int distance = Math.Abs(samples[sampleIndex] - centroids[centroidIndex]); + if (distance < bestDistance) + { + bestDistance = distance; + bestIndex = centroidIndex; + } + } + + indices[sampleIndex] = (byte)bestIndex; + distortion += (long)bestDistance * bestDistance; + } + + return distortion; + } +}