use pinned buffers in Vp8LHistogram

Author: antonfirsov

src/ImageSharp/Formats/Webp/Lossless/BackwardReferenceEncoder.cs

@@ -85,7 +85,7 @@ public static Vp8LBackwardRefs GetBackwardReferences(
             }
 
             // Keep the best backward references.
-            using Vp8LHistogram histo = new(memoryAllocator, worst, cacheBitsTmp);
+            using Vp8LHistogram histo = Vp8LHistogram.Create(memoryAllocator, worst, cacheBitsTmp);
             double bitCost = histo.EstimateBits(stats, bitsEntropy);
 
             if (lz77TypeBest == 0 || bitCost < bitCostBest)
@@ -102,7 +102,7 @@ public static Vp8LBackwardRefs GetBackwardReferences(
         {
             Vp8LHashChain hashChainTmp = lz77TypeBest == (int)Vp8LLz77Type.Lz77Standard ? hashChain : hashChainBox!;
             BackwardReferencesTraceBackwards(width, height, memoryAllocator, bgra, cacheBits, hashChainTmp, best, worst);
-            using Vp8LHistogram histo = new(memoryAllocator, worst, cacheBits);
+            using Vp8LHistogram histo = Vp8LHistogram.Create(memoryAllocator, worst, cacheBits);
             double bitCostTrace = histo.EstimateBits(stats, bitsEntropy);
             if (bitCostTrace < bitCostBest)
             {

src/ImageSharp/Formats/Webp/Lossless/CostModel.cs

@@ -37,7 +37,7 @@ public CostModel(MemoryAllocator memoryAllocator, int literalArraySize)
 
     public void Build(int xSize, int cacheBits, Vp8LBackwardRefs backwardRefs)
     {
-        using Vp8LHistogram histogram = new(this.memoryAllocator, cacheBits);
+        using Vp8LHistogram histogram = Vp8LHistogram.Create(this.memoryAllocator, cacheBits);
 
         // The following code is similar to HistogramCreate but converts the distance to plane code.
         for (int i = 0; i < backwardRefs.Refs.Count; i++)

src/ImageSharp/Formats/Webp/Lossless/Vp8LEncoder.cs

@@ -589,7 +589,7 @@ private void EncodeImage(int width, int height, bool useCache, CrunchConfig conf
             Vp8LBackwardRefs refsTmp = this.Refs[refsBest.Equals(this.Refs[0]) ? 1 : 0];
 
             this.bitWriter.Reset(bwInit);
-            using Vp8LHistogram tmpHisto = new(this.memoryAllocator, cacheBits);
+            using Vp8LHistogram tmpHisto = Vp8LHistogram.Create(this.memoryAllocator, cacheBits);
             using Vp8LHistogramSet histogramImage = new(this.memoryAllocator, histogramImageXySize, cacheBits);
 
             // Build histogram image and symbols from backward references.

src/ImageSharp/Formats/Webp/Lossless/Vp8LHistogram.cs

@@ -10,38 +10,27 @@
 
 namespace SixLabors.ImageSharp.Formats.Webp.Lossless;
 
-internal sealed class Vp8LHistogram : IDisposable
+internal sealed unsafe class Vp8LHistogram : IDisposable
 {
     private const uint NonTrivialSym = 0xffffffff;
-    private readonly IMemoryOwner<uint> buffer;
+    private readonly IMemoryOwner<uint>? bufferOwner;
+    private readonly Memory<uint> buffer;
+    private readonly MemoryHandle bufferHandle;
+
+    private readonly uint* red;
+    private readonly uint* blue;
+    private readonly uint* alpha;
+    private readonly uint* distance;
+    private readonly uint* literal;
+    private readonly uint* isUsed;
+
     private const int RedSize = WebpConstants.NumLiteralCodes;
     private const int BlueSize = WebpConstants.NumLiteralCodes;
     private const int AlphaSize = WebpConstants.NumLiteralCodes;
     private const int DistanceSize = WebpConstants.NumDistanceCodes;
     public const int LiteralSize = WebpConstants.NumLiteralCodes + WebpConstants.NumLengthCodes + (1 << WebpConstants.MaxColorCacheBits) + 1;
     private const int UsedSize = 5; // 5 for literal, red, blue, alpha, distance
     public const int BufferSize = RedSize + BlueSize + AlphaSize + DistanceSize + LiteralSize + UsedSize;
-    private readonly bool isSetMember;
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="Vp8LHistogram"/> class.
-    /// </summary>
-    /// <param name="memoryAllocator">The memory allocator.</param>
-    /// <param name="refs">The backward references to initialize the histogram with.</param>
-    /// <param name="paletteCodeBits">The palette code bits.</param>
-    public Vp8LHistogram(MemoryAllocator memoryAllocator, Vp8LBackwardRefs refs, int paletteCodeBits)
-        : this(memoryAllocator, paletteCodeBits) => this.StoreRefs(refs);
-
-    /// <summary>
-    /// Initializes a new instance of the <see cref="Vp8LHistogram"/> class.
-    /// </summary>
-    /// <param name="memoryAllocator">The memory allocator.</param>
-    /// <param name="paletteCodeBits">The palette code bits.</param>
-    public Vp8LHistogram(MemoryAllocator memoryAllocator, int paletteCodeBits)
-    {
-        this.buffer = memoryAllocator.Allocate<uint>(BufferSize, AllocationOptions.Clean);
-        this.PaletteCodeBits = paletteCodeBits;
-    }
 
     /// <summary>
     /// Initializes a new instance of the <see cref="Vp8LHistogram"/> class.
@@ -52,7 +41,7 @@ public Vp8LHistogram(MemoryAllocator memoryAllocator, int paletteCodeBits)
     /// <param name="buffer">The backing buffer.</param>
     /// <param name="refs">The backward references to initialize the histogram with.</param>
     /// <param name="paletteCodeBits">The palette code bits.</param>
-    public Vp8LHistogram(IMemoryOwner<uint> buffer, Vp8LBackwardRefs refs, int paletteCodeBits)
+    public Vp8LHistogram(Memory<uint> buffer, Vp8LBackwardRefs refs, int paletteCodeBits)
         : this(buffer, paletteCodeBits) => this.StoreRefs(refs);
 
     /// <summary>
@@ -63,11 +52,20 @@ public Vp8LHistogram(IMemoryOwner<uint> buffer, Vp8LBackwardRefs refs, int palet
     /// </remarks>
     /// <param name="buffer">The backing buffer.</param>
     /// <param name="paletteCodeBits">The palette code bits.</param>
-    public Vp8LHistogram(IMemoryOwner<uint> buffer, int paletteCodeBits)
+    /// <param name="bufferOwner">Optional buffer owner to dispose.</param>
+    public Vp8LHistogram(Memory<uint> buffer, int paletteCodeBits, IMemoryOwner<uint>? bufferOwner = null)
     {
+        this.bufferOwner = bufferOwner;
         this.buffer = buffer;
+        this.bufferHandle = this.buffer.Pin();
         this.PaletteCodeBits = paletteCodeBits;
-        this.isSetMember = true;
+
+        this.red = (uint*)this.bufferHandle.Pointer;
+        this.blue = this.red + RedSize;
+        this.alpha = this.blue + BlueSize;
+        this.distance = this.alpha + AlphaSize;
+        this.literal = this.distance + DistanceSize;
+        this.isUsed = this.literal + LiteralSize;
     }
 
     /// <summary>
@@ -95,22 +93,43 @@ public Vp8LHistogram(IMemoryOwner<uint> buffer, int paletteCodeBits)
     /// </summary>
     public double BlueCost { get; set; }
 
-    public Span<uint> Red => this.buffer.GetSpan()[..RedSize];
+    public Span<uint> Red => new(this.red, RedSize);
 
-    public Span<uint> Blue => this.buffer.GetSpan().Slice(RedSize, BlueSize);
+    public Span<uint> Blue => new(this.blue, BlueSize);
 
-    public Span<uint> Alpha => this.buffer.GetSpan().Slice(RedSize + BlueSize, AlphaSize);
+    public Span<uint> Alpha => new(this.alpha, AlphaSize);
 
-    public Span<uint> Distance => this.buffer.GetSpan().Slice(RedSize + BlueSize + AlphaSize, DistanceSize);
+    public Span<uint> Distance => new(this.distance, DistanceSize);
 
-    public Span<uint> Literal => this.buffer.GetSpan().Slice(RedSize + BlueSize + AlphaSize + DistanceSize, LiteralSize);
+    public Span<uint> Literal => new(this.literal, LiteralSize);
 
     public uint TrivialSymbol { get; set; }
 
-    private Span<uint> IsUsedSpan => this.buffer.GetSpan().Slice(RedSize + BlueSize + AlphaSize + DistanceSize + LiteralSize, UsedSize);
+    private Span<uint> IsUsedSpan => new(this.isUsed, UsedSize);
+
+    private Span<uint> TotalSpan => new(this.red, BufferSize);
 
     public bool IsDisposed { get; set; }
 
+    /// <summary>
+    /// Creates an <see cref="Vp8LHistogram"/> that is not a member of a <see cref="Vp8LHistogramSet"/>.
+    /// </summary>
+    public static Vp8LHistogram Create(MemoryAllocator memoryAllocator, int paletteCodeBits)
+    {
+        IMemoryOwner<uint> bufferOwner = memoryAllocator.Allocate<uint>(BufferSize, AllocationOptions.Clean);
+        return new Vp8LHistogram(bufferOwner.Memory, paletteCodeBits, bufferOwner);
+    }
+
+    /// <summary>
+    /// Creates an <see cref="Vp8LHistogram"/> that is not a member of a <see cref="Vp8LHistogramSet"/>.
+    /// </summary>
+    public static Vp8LHistogram Create(MemoryAllocator memoryAllocator, Vp8LBackwardRefs refs, int paletteCodeBits)
+    {
+        Vp8LHistogram histogram = Create(memoryAllocator, paletteCodeBits);
+        histogram.StoreRefs(refs);
+        return histogram;
+    }
+
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     public bool IsUsed(int index) => this.IsUsedSpan[index] == 1u;
 
@@ -140,7 +159,7 @@ public void CopyTo(Vp8LHistogram other)
 
     public void Clear()
     {
-        this.buffer.Clear();
+        this.TotalSpan.Clear();
         this.PaletteCodeBits = 0;
         this.BitCost = 0;
         this.LiteralCost = 0;
@@ -607,11 +626,8 @@ public void Dispose()
     {
         if (!this.IsDisposed)
         {
-            if (!this.isSetMember)
-            {
-                this.buffer.Dispose();
-            }
-
+            this.bufferHandle.Dispose();
+            this.bufferOwner?.Dispose();
             this.IsDisposed = true;
         }
     }

src/ImageSharp/Formats/Webp/Lossless/Vp8LHistogramSet.cs

@@ -23,8 +23,8 @@ public Vp8LHistogramSet(MemoryAllocator memoryAllocator, int capacity, int cache
         this.items = new List<Vp8LHistogram>(capacity);
         for (int i = 0; i < capacity; i++)
         {
-            SetItemMemoryOwner owner = new(this.buffer.Memory.Slice(Vp8LHistogram.BufferSize * i, Vp8LHistogram.BufferSize));
-            this.items.Add(new Vp8LHistogram(owner, cacheBits));
+            Memory<uint> subBuffer = this.buffer.Memory.Slice(Vp8LHistogram.BufferSize * i, Vp8LHistogram.BufferSize);
+            this.items.Add(new Vp8LHistogram(subBuffer, cacheBits));
         }
     }
 
@@ -35,8 +35,8 @@ public Vp8LHistogramSet(MemoryAllocator memoryAllocator, Vp8LBackwardRefs refs,
         this.items = new List<Vp8LHistogram>(capacity);
         for (int i = 0; i < capacity; i++)
         {
-            SetItemMemoryOwner owner = new(this.buffer.Memory.Slice(Vp8LHistogram.BufferSize * i, Vp8LHistogram.BufferSize));
-            this.items.Add(new Vp8LHistogram(owner, refs, cacheBits));
+            Memory<uint> subBuffer = this.buffer.Memory.Slice(Vp8LHistogram.BufferSize * i, Vp8LHistogram.BufferSize);
+            this.items.Add(new Vp8LHistogram(subBuffer, refs, cacheBits));
         }
     }
 
@@ -82,13 +82,13 @@ public void Dispose()
             return;
         }
 
-        this.buffer.Dispose();
-
         foreach (Vp8LHistogram item in this.items)
         {
+            // First, make sure to unpin individual sub buffers.
             item?.Dispose();
         }
 
+        this.buffer.Dispose();
         this.items.Clear();
         this.isDisposed = true;
     }
@@ -107,16 +107,4 @@ private void CheckDisposed()
     }
 
     private static void ThrowDisposed() => throw new ObjectDisposedException(nameof(Vp8LHistogramSet));
-
-    private sealed class SetItemMemoryOwner : IMemoryOwner<uint>
-    {
-        public SetItemMemoryOwner(Memory<uint> memory) => this.Memory = memory;
-
-        public Memory<uint> Memory { get; }
-
-        public void Dispose()
-        {
-            // Do nothing, the underlying memory is owned by the parent set.
-        }
-    }
 }

tests/ImageSharp.Tests/Formats/WebP/DominantCostRangeTests.cs

@@ -25,12 +25,10 @@ public void UpdateDominantCostRange_Works()
     {
         // arrange
         DominantCostRange dominantCostRange = new();
-        using Vp8LHistogram histogram = new(Configuration.Default.MemoryAllocator, 10)
-        {
-            LiteralCost = 1.0d,
-            RedCost = 2.0d,
-            BlueCost = 3.0d
-        };
+        using Vp8LHistogram histogram = Vp8LHistogram.Create(Configuration.Default.MemoryAllocator, 10);
+        histogram.LiteralCost = 1.0d;
+        histogram.RedCost = 2.0d;
+        histogram.BlueCost = 3.0d;
 
         // act
         dominantCostRange.UpdateDominantCostRange(histogram);
@@ -59,13 +57,12 @@ public void GetHistoBinIndex_Works(int partitions, int expectedIndex)
             RedMax = 191.0,
             RedMin = 109.0
         };
-        using Vp8LHistogram histogram = new(Configuration.Default.MemoryAllocator, 6)
-        {
-            LiteralCost = 247.0d,
-            RedCost = 112.0d,
-            BlueCost = 202.0d,
-            BitCost = 733.0d
-        };
+        using Vp8LHistogram histogram = Vp8LHistogram.Create(Configuration.Default.MemoryAllocator, 6);
+        histogram.LiteralCost = 247.0d;
+        histogram.RedCost = 112.0d;
+        histogram.BlueCost = 202.0d;
+        histogram.BitCost = 733.0d;
+
         dominantCostRange.UpdateDominantCostRange(histogram);
 
         // act

tests/ImageSharp.Tests/Formats/WebP/Vp8LHistogramTests.cs

@@ -78,15 +78,15 @@ private static void RunAddVectorTest()
         }
 
         MemoryAllocator memoryAllocator = Configuration.Default.MemoryAllocator;
-        using Vp8LHistogram histogram0 = new(memoryAllocator, backwardRefs, 3);
-        using Vp8LHistogram histogram1 = new(memoryAllocator, backwardRefs, 3);
+        using Vp8LHistogram histogram0 = Vp8LHistogram.Create(memoryAllocator, backwardRefs, 3);
+        using Vp8LHistogram histogram1 = Vp8LHistogram.Create(memoryAllocator, backwardRefs, 3);
         for (int i = 0; i < 5; i++)
         {
             histogram0.IsUsed(i, true);
             histogram1.IsUsed(i, true);
         }
 
-        using Vp8LHistogram output = new(memoryAllocator, 3);
+        using Vp8LHistogram output = Vp8LHistogram.Create(memoryAllocator, 3);
 
         // act
         histogram0.Add(histogram1, output);