Match current libaom entropy normalization

HEIF_IMPLEMENTATION_PLAN.md

@@ -826,7 +826,7 @@ Encoder verification contract:
 - [~] Finalized transform coefficients and packed EOB/type state now use raster-ordered, per-superblock plane segments matching current libaom's coefficient-pool geometry. One ImageSharp allocator owner replaces libaom's separate coefficient, EOB, and entropy-context allocations while preserving the full 1024 luma and 256-per-chroma 4x4 state capacity of a 128x128 4:2:0 superblock; the forward transform and mode-decision stages still need to populate this owner.
 - [~] Tile partition writing now follows current libaom's recursive `write_modes_sb` preorder traversal and `update_ext_partition_context` edge updates directly. The obsolete SVT-derived global geometry catalog and its unimplemented lookup are removed; transform geometry is derived in libaom's bounded 64x64 residual order, fixed intra transform-size symbols use the reference depth and neighbor contexts, frame-edge and segmentation syntax use mode-information units, and 128x128 CDEF units use libaom's 0-to-3 indexing and first-block strength ownership. Partition and mode analysis still need to populate these retained decisions; variable inter-transform syntax remains part of later inter-frame support.
 - [ ] Implement legal deblocking, CDEF, restoration, super-resolution, and film-grain signaling decisions.
-- [~] The coefficient symbol encoder now reuses tile-lifetime level and context workspaces instead of allocating per transform. The reference-type symbol encoder is passed normally through tile traversal, and the operation boundary owns the allocator-backed item payload stream for exactly one synchronous encode. Every remaining encoder fragment must be audited before it becomes active.
+- [~] The coefficient symbol encoder now reuses tile-lifetime level and context workspaces instead of allocating per transform. Its range coder matches current libaom's 64-bit coding window, bulk big-endian byte flush, and backward carry propagation while using one byte of allocator scratch per estimated output byte instead of the former 16-bit pre-carry storage. The reference-type symbol encoder is passed normally through tile traversal, and the operation boundary owns the allocator-backed item payload stream for exactly one synchronous encode. Every remaining encoder fragment must be audited before it becomes active.
 - [~] The planar conversion, forward transform, and forward quantizer use descending SIMD dispatch: Vector512, Vector256, Vector128, then scalar. Apply the same rule to every later hot-path family.
 - [~] Forward-quantizer FeatureTestRunner and zero-allocation tests compare every hardware tier with an independent scan-order scalar oracle shaped from current-main libaom. Both passed direct net11 VSTest in Release.
 - [~] The combined-frame writer now completes the byte-counted uncompressed frame header before starting the optional multi-tile tile-group flag, matching current libaom's separate frame-header and tile-group writers. A non-uniform two-tile round trip verifies the explicit boundaries, both tile payloads, and complete stream consumption through direct net11 VSTest in Release.

src/ImageSharp/Formats/Heif/Av1/Entropy/Av1SymbolWriter.cs

@@ -2,6 +2,7 @@
 // Licensed under the Six Labors Split License.
 
 using System.Buffers;
+using System.Buffers.Binary;
 using SixLabors.ImageSharp.Memory;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Entropy;
@@ -14,7 +15,7 @@ internal class Av1SymbolWriter : IDisposable
     /// <summary>
     /// The lower endpoint of the current coding interval.
     /// </summary>
-    private uint low;
+    private ulong low;
 
     /// <summary>
     /// The width of the current normalized coding interval.
@@ -35,9 +36,9 @@ internal class Av1SymbolWriter : IDisposable
     private readonly Configuration configuration;
 
     /// <summary>
-    /// The pre-carry output values accumulated during renormalization.
+    /// The output bytes accumulated during renormalization.
     /// </summary>
-    private readonly AutoExpandingMemory<ushort> memory;
+    private readonly AutoExpandingMemory<byte> memory;
 
     /// <summary>
     /// Indicates whether encoded symbols adapt their distributions.
@@ -45,7 +46,7 @@ internal class Av1SymbolWriter : IDisposable
     private readonly bool updateCdf;
 
     /// <summary>
-    /// The next pre-carry output position.
+    /// The next output byte position.
     /// </summary>
     private int position;
 
@@ -58,7 +59,7 @@ internal class Av1SymbolWriter : IDisposable
     public Av1SymbolWriter(Configuration configuration, int initialSize, bool updateCdf = true)
     {
         this.configuration = configuration;
-        this.memory = new AutoExpandingMemory<ushort>(configuration, (initialSize + 1) >> 1);
+        this.memory = new AutoExpandingMemory<byte>(configuration, initialSize);
         this.updateCdf = updateCdf;
     }
 
@@ -129,22 +130,29 @@ internal class Av1SymbolWriter : IDisposable
     /// <returns>An owner containing the shortest byte sequence that preserves every encoded symbol.</returns>
     public IMemoryOwner<byte> Exit()
     {
-        // Round the low endpoint into the current interval so the emitted prefix selects every symbol encoded so far
+        // Round the low endpoint into the current interval so the emitted prefix selects every symbol encoded so far,
         // regardless of the bits that follow it.
-        uint l = this.low;
+        ulong l = this.low;
         int c = this.cnt;
         int pos = this.position;
         int s = 10;
-        uint m = 0x3FFFU;
+        ulong m = 0x3FFFU;
-        uint e = ((l + m) & ~m) | (m + 1);
+        ulong e = ((l + m) & ~m) | (m + 1);
         s += c;
-        Span<ushort> buffer = this.memory.GetSpan(this.position + ((s + 7) >> 3));
+        int pendingByteCount = Math.Max((s + 7) >> 3, 0);
+        Span<byte> buffer = this.memory.GetSpan(pos + pendingByteCount);
         if (s > 0)
         {
-            uint n = (1U << (c + 16)) - 1;
+            ulong n = (1UL << (c + 16)) - 1;
             do
             {
-                buffer[pos] = (ushort)(e >> (c + 16));
+                ushort value = (ushort)(e >> (c + 16));
+                buffer[pos] = (byte)value;
+                if ((value & 0x100) != 0)
+                {
+                    PropagateCarryBackward(buffer, pos - 1);
+                }
+
                 pos++;
                 e &= n;
                 s -= 8;
@@ -154,20 +162,8 @@ internal class Av1SymbolWriter : IDisposable
             while (s > 0);
         }
 
-        c = Math.Max((s + 7) >> 3, 0);
+        IMemoryOwner<byte> output = this.configuration.MemoryAllocator.Allocate<byte>(pos);
-        IMemoryOwner<byte> output = this.configuration.MemoryAllocator.Allocate<byte>(pos + c);
+        buffer[..pos].CopyTo(output.GetSpan()[..pos]);
-
-        // Pre-carry values use 16-bit elements so a byte plus a propagated carry can coexist. Walking backwards folds
-        // each carry into the preceding byte without shifting the buffered sequence.
-        Span<byte> outputSlice = output.GetSpan()[(output.Length() - pos)..];
-        c = 0;
-        while (pos > 0)
-        {
-            pos--;
-            c = buffer[pos] + c;
-            outputSlice[pos] = (byte)c;
-            c >>= 8;
-        }
 
         return output;
     }
@@ -179,7 +175,7 @@ internal class Av1SymbolWriter : IDisposable
     /// <param name="frequency">The probability that the value is true, scaled by 32768.</param>
     private void EncodeBoolQ15(bool val, uint frequency)
     {
-        uint l;
+        ulong l;
         uint r;
         uint v;
         DebugGuard.MustBeGreaterThan(frequency, 0U, nameof(frequency));
@@ -227,7 +223,7 @@ internal class Av1SymbolWriter : IDisposable
     private void EncodeIntegerQ15(uint lowFrequency, uint highFrequency, int symbol, int numberOfSymbols)
     {
         const int totalShift = 7 - Av1Distribution.ProbabilityShift - Av1Distribution.CdfShift;
-        uint l = this.low;
+        ulong l = this.low;
         uint r = this.rng;
         DebugGuard.MustBeLessThanOrEqualTo(32768U, r, nameof(r));
         DebugGuard.MustBeLessThanOrEqualTo(highFrequency, lowFrequency, nameof(highFrequency));
@@ -261,42 +257,61 @@ internal class Av1SymbolWriter : IDisposable
     /// </summary>
     /// <param name="low">The new value of <see cref="low"/>.</param>
     /// <param name="rng">The new value of <see cref="rng"/>.</param>
-    private void Normalize(uint low, uint rng)
+    private void Normalize(ulong low, uint rng)
     {
-        int d;
+        int c = this.cnt;
-        int c;
-        int s;
-        c = this.cnt;
         DebugGuard.MustBeLessThanOrEqualTo(rng, 65535U, nameof(rng));
-        d = 15 - Av1Math.MostSignificantBit(rng);
+        int d = 15 - Av1Math.MostSignificantBit(rng);
-        s = c + d;
+        int s = c + d;
 
-        // The 32-bit low endpoint is flushed whenever a byte becomes available. Retaining pre-carry values as
+        // Keeping 16 bits free for the next symbol allows the 64-bit coding window to flush up to eight completed
-        // ushort elements defers carry propagation until Exit without requiring a separate wider coding window.
+        // bytes together while preserving one carry bit.
-        if (s >= 0)
+        if (s >= 40)
         {
-            uint m;
+            Span<byte> buffer = this.memory.GetSpan(this.position + sizeof(ulong));
-            Span<ushort> buffer = this.memory.GetSpan(this.position + 2);
+            int readyByteCount = (s >> 3) + 1;
-
+            c += 24 - (readyByteCount << 3);
-            c += 16;
+            ulong output = low >> c;
-            m = (1U << c) - 1;
+            low &= (1UL << c) - 1;
-            if (s >= 8)
+            ulong carryMask = 1UL << (readyByteCount << 3);
+            bool hasCarry = (output & carryMask) != 0;
+            output &= carryMask - 1;
+
+            // Writing one big-endian word avoids a byte-at-a-time hot loop. Only readyByteCount bytes become part
+            // of the logical output; the following bytes are overwritten by the next flush.
+            BinaryPrimitives.WriteUInt64BigEndian(
+                buffer.Slice(this.position, sizeof(ulong)),
+                output << ((sizeof(ulong) - readyByteCount) << 3));
+
+            if (hasCarry)
             {
-                buffer[this.position] = (ushort)(low >> c);
+                PropagateCarryBackward(buffer, this.position - 1);
-                this.position++;
-                low &= m;
-                c -= 8;
-                m >>= 8;
             }
 
-            buffer[this.position] = (ushort)(low >> c);
+            this.position += readyByteCount;
-            this.position++;
             s = c + d - 24;
-            low &= m;
         }
 
         this.low = low << d;
         this.rng = rng << d;
         this.cnt = s;
     }
+
+    /// <summary>
+    /// Adds a carry to the completed output prefix.
+    /// </summary>
+    /// <param name="buffer">The accumulated output bytes.</param>
+    /// <param name="offset">The final completed byte.</param>
+    private static void PropagateCarryBackward(Span<byte> buffer, int offset)
+    {
+        int carry;
+        do
+        {
+            int sum = buffer[offset] + 1;
+            buffer[offset] = (byte)sum;
+            carry = sum >> 8;
+            offset--;
+        }
+        while (carry != 0);
+    }
 }

tests/ImageSharp.Tests/Formats/Heif/Av1/Av1EntropyTests.cs

@@ -9,6 +9,7 @@ using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
 using SixLabors.ImageSharp.Memory;
+using SixLabors.ImageSharp.Tests.Memory;
 
 namespace SixLabors.ImageSharp.Tests.Formats.Heif.Av1;
 
@@ -17,6 +18,43 @@ public class Av1EntropyTests
 {
     private const int BaseQIndex = 23;
 
+    [Fact]
+    public void SymbolWriterMatchesCurrentLibaomCarryRegression()
+    {
+        using Av1SymbolWriter writer = new(Configuration.Default, 1, updateCdf: false);
+        writer.WriteBoolean(false, 16_384);
+        writer.WriteBoolean(false, 16_384);
+        writer.WriteBoolean(true, 512);
+        writer.WriteBoolean(false, 8_192);
+        using IMemoryOwner<byte> encoded = writer.Exit();
+
+        Assert.Equal(2, encoded.Memory.Length);
+        Assert.Equal(63, encoded.Memory.Span[0]);
+    }
+
+    [Fact]
+    public void SymbolWriterUsesOneByteOfScratchPerEstimatedOutputByte()
+    {
+        const int initialSize = 257;
+        TestMemoryAllocator allocator = new();
+        allocator.EnableNonThreadSafeLogging();
+        Configuration configuration = Configuration.Default.Clone();
+        configuration.MemoryAllocator = allocator;
+        TestMemoryAllocator.AllocationRequest allocation;
+
+        using (Av1SymbolWriter writer = new(configuration, initialSize, updateCdf: false))
+        {
+            writer.WriteLiteral(false);
+            allocation = Assert.Single(allocator.AllocationLog);
+
+            Assert.Equal(typeof(byte), allocation.ElementType);
+            Assert.Equal(initialSize, allocation.Length);
+        }
+
+        TestMemoryAllocator.ReturnRequest returned = Assert.Single(allocator.ReturnLog);
+        Assert.Equal(allocation.HashCodeOfBuffer, returned.HashCodeOfBuffer);
+    }
+
     [Fact]
     public void ReadRandomLiteral()
     {