Consolidate AV1 encoder block scratch

HEIF_IMPLEMENTATION_PLAN.md

@@ -818,7 +818,7 @@ Encoder verification contract:
 ### 6. Build the complete AV1 frame encoder
 
 - [~] SIMD-first RGB-to-native-plane conversion now feeds eight-bit and high-bit-depth bordered AV1 source frames directly, preserving ImageSharp's arbitrary packed-pixel input contract without an intermediate full-frame native-plane copy.
-- [~] Forward transform families, transform workspace, and a zero-allocation transform-then-quantize block boundary exist locally. The block boundary matches current libaom's selected full-transform, quantization, qcoeff, dqcoeff, EOB, and transform-type data flow; frame traversal still needs to select blocks and supply coefficient-owner slices.
+- [~] Forward transform families, transform workspace, and a zero-allocation transform-then-quantize block boundary exist locally. The block boundary matches current libaom's selected full-transform, quantization, qcoeff, dqcoeff, EOB, and transform-type data flow. One reusable 61 KiB allocator owner supplies tightly packed residual, aligned transform-coefficient, dequantized-coefficient, and transform scratch spans across transform blocks; quantized coefficients write directly to the retained frame coefficient owner instead of being duplicated. The exact owner type, length, span partitioning, and exactly-once return are verified, and the adjacent residual, transform, quantizer, and block tests pass 12 of 12 through direct net11 VSTest in Release. Frame traversal still needs to select blocks and supply coefficient-owner slices.
 - [~] Symbol writer, coefficient writer, and tile writer fragments exist locally.
 - [~] A non-owning encoder-frame view now separates visible conversion regions from coded regions and performs complete left, top, right, bottom, and corner extension across each bordered plane. Current libaom uses 8-sample-aligned coded dimensions, a 32-sample-aligned luma stride with chroma stride derived from it, and a 64-pixel luma border for non-resized all-intra encoding. The frame-encoder boundary now converts packed pixels directly into those final source planes before extension; the containing encode operation still needs to connect matching source and reconstruction plane rents with ordinary `using` lifetimes.
 - [~] Temporal delimiter, sequence header, frame header, and combined-frame tile-group writing exist locally. The remaining required metadata, padding, and encoder-wide syntax paths are not complete.

src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1EncoderBlockWorkspace.cs

@@ -0,0 +1,76 @@
+// Copyright (c) Six Labors.
+// Licensed under the Six Labors Split License.
+
+using System.Buffers;
+using System.Runtime.InteropServices;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
+using SixLabors.ImageSharp.Memory;
+
+namespace SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline;
+
+/// <summary>
+/// Owns the reusable spatial, transform, and reconstruction storage for AV1 block encoding.
+/// </summary>
+internal sealed class Av1EncoderBlockWorkspace : IDisposable
+{
+    /// <summary>
+    /// The maximum number of spatial residual samples in one AV1 transform block.
+    /// </summary>
+    public const int MaximumResidualCount = Av1Constants.MaxTransformSize * Av1Constants.MaxTransformSize;
+
+    /// <summary>
+    /// The maximum number of coded coefficients after AV1 removes the uncoded half of 64-point axes.
+    /// </summary>
+    public const int MaximumCoefficientCount = (Av1Constants.MaxTransformSize / 2) * (Av1Constants.MaxTransformSize / 2);
+
+    /// <summary>
+    /// The complete workspace length in signed-integer storage elements.
+    /// </summary>
+    public const int StorageLength = ResidualStorageLength + MaximumCoefficientCount + MaximumCoefficientCount + Av1TransformWorkspace.MaximumLength;
+
+    private const int ResidualStorageLength = MaximumResidualCount / 2;
+    private const int TransformCoefficientOffset = ResidualStorageLength;
+    private const int DequantizedCoefficientOffset = TransformCoefficientOffset + MaximumCoefficientCount;
+    private const int TransformWorkspaceOffset = DequantizedCoefficientOffset + MaximumCoefficientCount;
+
+    /// <summary>
+    /// Owns the complete reusable block workspace in 32-bit elements so every transform region is naturally aligned.
+    /// </summary>
+    private readonly IMemoryOwner<int> owner;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="Av1EncoderBlockWorkspace"/> class.
+    /// </summary>
+    /// <param name="configuration">The configuration providing the encoder allocator.</param>
+    public Av1EncoderBlockWorkspace(Configuration configuration)
+        => this.owner = configuration.MemoryAllocator.Allocate<int>(StorageLength);
+
+    /// <summary>
+    /// Gets the maximum-size spatial residual workspace as a compact 16-bit view of the aligned owner.
+    /// </summary>
+    public Span<short> Residual
+        => MemoryMarshal.Cast<int, short>(this.owner.Memory.Span[..ResidualStorageLength]);
+
+    /// <summary>
+    /// Gets the maximum-size forward-transform coefficient workspace.
+    /// </summary>
+    public Span<int> TransformCoefficients
+        => this.owner.Memory.Span.Slice(TransformCoefficientOffset, MaximumCoefficientCount);
+
+    /// <summary>
+    /// Gets the maximum-size dequantized reconstruction coefficient workspace.
+    /// </summary>
+    public Span<int> DequantizedCoefficients
+        => this.owner.Memory.Span.Slice(DequantizedCoefficientOffset, MaximumCoefficientCount);
+
+    /// <summary>
+    /// Gets the reusable two-dimensional transform workspace.
+    /// </summary>
+    public Span<int> TransformWorkspace
+        => this.owner.Memory.Span.Slice(TransformWorkspaceOffset, Av1TransformWorkspace.MaximumLength);
+
+    /// <summary>
+    /// Releases the reusable block workspace.
+    /// </summary>
+    public void Dispose() => this.owner.Dispose();
+}

src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1TransformBlockEncoder.cs

@@ -15,12 +15,8 @@ internal static class Av1TransformBlockEncoder
     /// <summary>
     /// Applies a lossy forward transform and quantization to one residual block.
     /// </summary>
-    /// <param name="residual">The spatial residual samples.</param>
-    /// <param name="residualStride">The number of residual samples between rows.</param>
-    /// <param name="transformCoefficients">The reusable forward-transform coefficient workspace.</param>
+    /// <param name="workspace">The reusable residual, transform, and reconstruction storage.</param>
     /// <param name="quantizedCoefficients">The retained entropy-coding coefficients.</param>
-    /// <param name="dequantizedCoefficients">The reusable reconstruction coefficients.</param>
-    /// <param name="transformWorkspace">The reusable internal transform workspace.</param>
     /// <param name="transformSize">The selected transform dimensions.</param>
     /// <param name="transformType">The selected compound transform type.</param>
     /// <param name="qIndex">The segment quantizer index.</param>
@@ -29,12 +25,8 @@ internal static class Av1TransformBlockEncoder
     /// <param name="bitDepth">The coded sample bit depth.</param>
     /// <param name="state">The retained transform type and end-of-block syntax.</param>
     public static void EncodeLossy(
-        Span<short> residual,
-        uint residualStride,
-        Span<int> transformCoefficients,
+        Av1EncoderBlockWorkspace workspace,
         Span<int> quantizedCoefficients,
-        Span<int> dequantizedCoefficients,
-        Span<int> transformWorkspace,
         Av1TransformSize transformSize,
         Av1TransformType transformType,
         int qIndex,
@@ -44,20 +36,20 @@ internal static class Av1TransformBlockEncoder
         ref Av1EncoderTransformBlockState state)
     {
         int coefficientCount = transformSize.GetAdjusted().GetSize2d();
-        Span<int> transformed = transformCoefficients[..coefficientCount];
+        Span<int> transformed = workspace.TransformCoefficients[..coefficientCount];
         Span<int> quantized = quantizedCoefficients[..coefficientCount];
-        Span<int> dequantized = dequantizedCoefficients[..coefficientCount];
+        Span<int> dequantized = workspace.DequantizedCoefficients[..coefficientCount];
 
         // The encoder keeps transformed, quantized, and reconstructed coefficients separate because mode decision
         // consumes all three while only the quantized values survive in the frame coefficient owner.
         Av1ForwardTransformer.Transform2d(
-            residual,
+            workspace.Residual,
             transformed,
-            residualStride,
+            (uint)transformSize.GetWidth(),
             transformType,
             transformSize,
             bitDepth.GetBitCount(),
-            transformWorkspace);
+            workspace.TransformWorkspace);
 
         state.EndOfBlock = Av1ForwardQuantizer.QuantizeLossy(
             transformed,

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

@@ -6,6 +6,7 @@ using SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline.Quantizers;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
+using SixLabors.ImageSharp.Tests.Memory;
 
 namespace SixLabors.ImageSharp.Tests.Formats.Heif.Av1;
 
@@ -37,21 +38,14 @@ public class Av1TransformBlockEncoderTests
         int width = transformSize.GetWidth();
         int height = transformSize.GetHeight();
         int coefficientCount = transformSize.GetAdjusted().GetSize2d();
-        short[] residual = new short[width * height];
-        int[] transformed = new int[coefficientCount];
         int[] quantized = new int[coefficientCount];
-        int[] dequantized = new int[coefficientCount];
-        int[] workspace = new int[Av1TransformWorkspace.MaximumLength];
-        FillResidual(residual, width, width, height, 4095);
+        using Av1EncoderBlockWorkspace workspace = new(Configuration.Default);
+        FillResidual(workspace.Residual, width, height, 4095);
         Av1EncoderTransformBlockState state = default;
 
         Av1TransformBlockEncoder.EncodeLossy(
-            residual,
-            (uint)width,
-            transformed,
-            quantized,
-            dequantized,
             workspace,
+            quantized,
             transformSize,
             Av1TransformType.DctDct,
             73,
@@ -64,12 +58,8 @@ public class Av1TransformBlockEncoderTests
         for (int iteration = 0; iteration < 16; iteration++)
         {
             Av1TransformBlockEncoder.EncodeLossy(
-                residual,
-                (uint)width,
-                transformed,
-                quantized,
-                dequantized,
                 workspace,
+                quantized,
                 transformSize,
                 Av1TransformType.DctDct,
                 73,
@@ -82,6 +72,34 @@ public class Av1TransformBlockEncoderTests
         Assert.Equal(0, GC.GetAllocatedBytesForCurrentThread() - before);
     }
 
+    /// <summary>
+    /// Verifies that the block workspace uses one exact-size allocator owner and returns it exactly once.
+    /// </summary>
+    [Fact]
+    public void BlockWorkspaceUsesOneExactSizeOwner()
+    {
+        TestMemoryAllocator allocator = new();
+        allocator.EnableNonThreadSafeLogging();
+        Configuration configuration = Configuration.Default.Clone();
+        configuration.MemoryAllocator = allocator;
+
+        TestMemoryAllocator.AllocationRequest allocation;
+        using (Av1EncoderBlockWorkspace workspace = new(configuration))
+        {
+            allocation = Assert.Single(allocator.AllocationLog);
+            Assert.Empty(allocator.ReturnLog);
+            Assert.Equal(typeof(int), allocation.ElementType);
+            Assert.Equal(Av1EncoderBlockWorkspace.StorageLength, allocation.Length);
+            Assert.Equal(Av1EncoderBlockWorkspace.MaximumResidualCount, workspace.Residual.Length);
+            Assert.Equal(Av1EncoderBlockWorkspace.MaximumCoefficientCount, workspace.TransformCoefficients.Length);
+            Assert.Equal(Av1EncoderBlockWorkspace.MaximumCoefficientCount, workspace.DequantizedCoefficients.Length);
+            Assert.Equal(Av1TransformWorkspace.MaximumLength, workspace.TransformWorkspace.Length);
+        }
+
+        TestMemoryAllocator.ReturnRequest returned = Assert.Single(allocator.ReturnLog);
+        Assert.Equal(allocation.AllocationId, returned.AllocationId);
+    }
+
     private static void ValidateBlock(
         Av1TransformSize transformSize,
         Av1TransformType transformType,
@@ -90,30 +108,24 @@ public class Av1TransformBlockEncoderTests
     {
         int width = transformSize.GetWidth();
         int height = transformSize.GetHeight();
-        int residualStride = width + 3;
         int coefficientCount = transformSize.GetAdjusted().GetSize2d();
         int sampleMaximum = (1 << bitDepth.GetBitCount()) - 1;
-        short[] residual = new short[residualStride * height];
         int[] expectedTransformed = new int[coefficientCount + 7];
         int[] expectedQuantized = new int[coefficientCount + 7];
         int[] expectedDequantized = new int[coefficientCount + 7];
-        int[] actualTransformed = new int[coefficientCount + 7];
         int[] actualQuantized = new int[coefficientCount + 7];
-        int[] actualDequantized = new int[coefficientCount + 7];
         int[] expectedWorkspace = new int[Av1TransformWorkspace.MaximumLength];
-        int[] actualWorkspace = new int[Av1TransformWorkspace.MaximumLength];
+        using Av1EncoderBlockWorkspace blockWorkspace = new(Configuration.Default);
         Array.Fill(expectedTransformed, int.MinValue);
         Array.Fill(expectedQuantized, int.MinValue);
         Array.Fill(expectedDequantized, int.MinValue);
-        Array.Fill(actualTransformed, int.MinValue);
         Array.Fill(actualQuantized, int.MinValue);
-        Array.Fill(actualDequantized, int.MinValue);
-        FillResidual(residual, residualStride, width, height, sampleMaximum);
+        FillResidual(blockWorkspace.Residual, width, height, sampleMaximum);
 
         Av1ForwardTransformer.Transform2d(
-            residual,
+            blockWorkspace.Residual,
             expectedTransformed.AsSpan(0, coefficientCount),
-            (uint)residualStride,
+            (uint)width,
             transformType,
             transformSize,
             bitDepth.GetBitCount(),
@@ -132,12 +144,8 @@ public class Av1TransformBlockEncoderTests
 
         Av1EncoderTransformBlockState actualState = default;
         Av1TransformBlockEncoder.EncodeLossy(
-            residual,
-            (uint)residualStride,
-            actualTransformed,
+            blockWorkspace,
             actualQuantized,
-            actualDequantized,
-            actualWorkspace,
             transformSize,
             transformType,
             qIndex,
@@ -146,21 +154,21 @@ public class Av1TransformBlockEncoderTests
             bitDepth,
             ref actualState);
 
-        Assert.Equal(expectedTransformed, actualTransformed);
+        AssertEqual(expectedTransformed, blockWorkspace.TransformCoefficients, coefficientCount);
         Assert.Equal(expectedQuantized, actualQuantized);
-        Assert.Equal(expectedDequantized, actualDequantized);
+        AssertEqual(expectedDequantized, blockWorkspace.DequantizedCoefficients, coefficientCount);
         Assert.Equal(expectedEndOfBlock, actualState.EndOfBlock);
         Assert.Equal(transformType, actualState.TransformType);
     }
 
-    private static void FillResidual(Span<short> residual, int stride, int width, int height, int sampleMaximum)
+    private static void FillResidual(Span<short> residual, int width, int height, int sampleMaximum)
     {
         for (int y = 0; y < height; y++)
         {
             for (int x = 0; x < width; x++)
             {
                 int index = (y * width) + x;
-                residual[(y * stride) + x] = (short)((index & 3) switch
+                residual[index] = (short)((index & 3) switch
                 {
                     0 => sampleMaximum,
                     1 => -sampleMaximum,
@@ -170,4 +178,12 @@ public class Av1TransformBlockEncoderTests
             }
         }
     }
+
+    private static void AssertEqual(ReadOnlySpan<int> expected, ReadOnlySpan<int> actual, int count)
+    {
+        for (int i = 0; i < count; i++)
+        {
+            Assert.Equal(expected[i], actual[i]);
+        }
+    }
 }