Implement recursive AV1 16x16 partition search

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

@@ -3,6 +3,7 @@
 
 using System.Buffers;
 using System.Runtime.InteropServices;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
 using SixLabors.ImageSharp.Memory;
 
@@ -31,7 +32,7 @@ internal sealed class Av1EncoderBlockWorkspace : IDisposable
         MaximumCoefficientCount +
         MaximumCoefficientCount +
         Av1TransformWorkspace.MaximumLength +
-        IntraBlockCopyStorageLength +
+        SharedModeDecisionStorageLength +
         PartitionContextStorageLength;
 
     private const int ResidualStorageLength = MaximumResidualCount / 2;
@@ -65,10 +66,31 @@ internal sealed class Av1EncoderBlockWorkspace : IDisposable
         IntraBlockCopyResidualStorageLength +
         IntraBlockCopyCoefficientStorageLength;
 
+    private const int ModeDecisionStorageLength = Av1EncoderModeDecisionWorkspace<ushort>.StorageLength;
+    private const int SharedModeDecisionStorageLength = ModeDecisionStorageLength > IntraBlockCopyStorageLength
+        ? ModeDecisionStorageLength
+        : IntraBlockCopyStorageLength;
+
     private const int PartitionContextStorageOffset =
-        IntraBlockCopySampleStorageOffset + IntraBlockCopyStorageLength;
+        IntraBlockCopySampleStorageOffset + SharedModeDecisionStorageLength;
+
+    private const int MaximumPartitionEdgeUnitCount =
+        2 * (1 << (Av1Constants.MaxSuperBlockSizeLog2 - Av1Constants.ModeInfoSizeLog2));
+
+    private const int PartitionContextBytesPerEdgeUnit =
+        Av1PartitionContext.StorageSize + (4 * sizeof(byte)) + Av1EncoderPaletteInfo.StorageSize;
+
+    private const int PartitionContextSlotByteLength =
+        MaximumPartitionEdgeUnitCount * PartitionContextBytesPerEdgeUnit;
+
+    private const int PartitionContextSlotLength =
+        PartitionContextSlotByteLength / sizeof(int);
 
-    private const int PartitionContextStorageLength = 4;
+    private const int PartitionTrialLevelCount =
+        Av1Constants.MaxSuperBlockSizeLog2 - 3 + 1;
+
+    private const int PartitionContextStorageLength =
+        PartitionContextSlotLength * PartitionTrialLevelCount;
 
     /// <summary>
     /// Owns the complete reusable block workspace in 32-bit elements so every transform region is naturally aligned.
@@ -107,13 +129,20 @@ internal sealed class Av1EncoderBlockWorkspace : IDisposable
         => this.owner.Memory.Span.Slice(TransformWorkspaceOffset, Av1TransformWorkspace.MaximumLength);
 
     /// <summary>
-    /// Gets storage for the coefficient and transform edges restored after an 8x8 partition trial.
+    /// Gets the disjoint edge snapshot used to restore one square partition-search level.
     /// </summary>
-    public Span<byte> PartitionContexts
-        => MemoryMarshal.AsBytes(
-            this.owner.Memory.Span.Slice(
-                PartitionContextStorageOffset,
-                PartitionContextStorageLength));
+    /// <param name="blockSize">The square partition node being evaluated.</param>
+    /// <returns>The maximum-size byte view reserved for that node depth.</returns>
+    public Span<byte> GetPartitionContextStorage(Av1BlockSize blockSize)
+    {
+        int blockSizeLog2 = Av1Math.Log2(blockSize.GetWidth());
+        int slotIndex = Av1Constants.MaxSuperBlockSizeLog2 - blockSizeLog2;
+        Span<int> storage = this.owner.Memory.Span.Slice(
+            PartitionContextStorageOffset + (slotIndex * PartitionContextSlotLength),
+            PartitionContextSlotLength);
+
+        return MemoryMarshal.AsBytes(storage);
+    }
 
     /// <summary>
     /// Gets the reusable storage used while comparing spatial, chroma-from-luma, filter-intra, and palette candidates.
@@ -127,7 +156,7 @@ internal sealed class Av1EncoderBlockWorkspace : IDisposable
         // Both phases can therefore reuse this aligned region without extending the owner or preserving stale scratch.
         Span<int> storage = this.owner.Memory.Span.Slice(
             IntraBlockCopySampleStorageOffset,
-            IntraBlockCopyStorageLength);
+            SharedModeDecisionStorageLength);
 
         return new(storage[..Av1EncoderModeDecisionWorkspace<TSample>.StorageLength]);
     }

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

@@ -16,9 +16,14 @@ internal readonly ref struct Av1EncoderModeDecisionWorkspace<TSample>
     where TSample : unmanaged
 {
     /// <summary>
-    /// The maximum number of samples in the encoder's fixed 8x8 transform block.
+    /// The largest coding-block dimension evaluated directly by the current partition search.
     /// </summary>
-    public const int MaximumSampleCount = 8 * 8;
+    public const int MaximumBlockDimension = 16;
+
+    /// <summary>
+    /// The maximum number of samples in one directly evaluated coding block.
+    /// </summary>
+    public const int MaximumSampleCount = MaximumBlockDimension * MaximumBlockDimension;
 
     /// <summary>
     /// The number of 4x4 transform blocks covering one 8x8 coding block.
@@ -30,7 +35,7 @@ internal readonly ref struct Av1EncoderModeDecisionWorkspace<TSample>
     /// </summary>
     public const int StorageLength = TransientStorageOffset + Av1EncoderPaletteWorkspace<ushort>.StorageLength;
 
-    private const int ReferenceBufferLength = 17;
+    private const int ReferenceBufferLength = (2 * MaximumBlockDimension) + 1;
     private const int ReferenceBufferCount = 4;
     private const int ReferenceStorageLength = ReferenceBufferCount * ReferenceBufferLength * sizeof(ushort) / sizeof(int);
     private const int CandidateSampleStorageOffset = ReferenceStorageLength;
@@ -40,9 +45,13 @@ internal readonly ref struct Av1EncoderModeDecisionWorkspace<TSample>
     private const int CandidateTransformBlockStorageOffset = CandidateCoefficientStorageOffset + CandidateCoefficientStorageLength;
     private const int CandidateTransformBlockStorageLength = CandidateTransformBlockCount;
     private const int TransformContextStorageOffset = CandidateTransformBlockStorageOffset + CandidateTransformBlockStorageLength;
-    private const int TransformContextStorageLength = 1;
+    private const int TransformContextStorageLength =
+        2 * (MaximumBlockDimension >> Av1Constants.ModeInfoSizeLog2) * sizeof(byte) / sizeof(int);
+
     private const int TransientStorageOffset = TransformContextStorageOffset + TransformContextStorageLength;
-    private const int ChromaFromLumaSampleCount = Av1ChromaFromLumaContext.BufferLine * 8;
+    private const int ChromaFromLumaSampleCount =
+        Av1ChromaFromLumaContext.BufferLine * MaximumBlockDimension;
+
     private const int ChromaFromLumaSampleStorageLength = ChromaFromLumaSampleCount * sizeof(short) / sizeof(int);
     private const int ChromaFromLumaBlueRateOffset = ChromaFromLumaSampleStorageLength;
     private const int ChromaFromLumaRedRateOffset = ChromaFromLumaBlueRateOffset + Av1ChromaFromLumaMath.AlphaCandidateCount;

src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.ModeDecision.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
+using System.Runtime.InteropServices;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Entropy;
 using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction;
@@ -40,6 +41,23 @@ internal static partial class Av1IntraSuperblockEncoder
     /// </summary>
     private static ReadOnlySpan<sbyte> AngleDeltaSearchOrder => [-3, -2, -1, 1, 2, 3];
 
+    /// <summary>
+    /// Gets partition candidates in the evaluation order used by the reference encoder.
+    /// </summary>
+    private static ReadOnlySpan<Av1PartitionType> PartitionSearchOrder =>
+    [
+        Av1PartitionType.None,
+        Av1PartitionType.Split,
+        Av1PartitionType.Horizontal,
+        Av1PartitionType.Vertical,
+        Av1PartitionType.HorizontalA,
+        Av1PartitionType.HorizontalB,
+        Av1PartitionType.VerticalA,
+        Av1PartitionType.VerticalB,
+        Av1PartitionType.Horizontal4,
+        Av1PartitionType.Vertical4
+    ];
+
     /// <summary>
     /// Builds the fixed 8x8 partition skeleton consumed by interleaved mode decision and tile writing.
     /// </summary>
@@ -165,32 +183,24 @@ internal static partial class Av1IntraSuperblockEncoder
             Av1BlockSize blockSize,
             Av1PartitionType preparedPartition)
         {
-            if (blockSize != Av1BlockSize.Block8x8)
+            if (blockSize is not Av1BlockSize.Block8x8 and not Av1BlockSize.Block16x16)
             {
                 return preparedPartition;
             }
 
             Point modeInfoPosition = blockOrigin >> Av1Constants.ModeInfoSizeLog2;
-            bool hasRows = modeInfoPosition.Y + 1 < this.picture.Parent.Common.ModeInfoRowCount;
-            bool hasColumns = modeInfoPosition.X + 1 < this.picture.Parent.Common.ModeInfoColumnCount;
-            if (!hasRows || !hasColumns)
-            {
-                // A clipped 8x8 node must split because its missing half cannot be represented by PARTITION_NONE.
-                for (int childIndex = 0; childIndex < 4; childIndex++)
-                {
-                    Point childOrigin = blockOrigin + new Size(
-                        (childIndex & 1) << Av1Constants.ModeInfoSizeLog2,
-                        (childIndex >> 1) << Av1Constants.ModeInfoSizeLog2);
+            bool hasRows =
+                modeInfoPosition.Y + blockSize.Get4x4HighCount() <= this.picture.Parent.Common.ModeInfoRowCount;
 
-                    Point childPosition = childOrigin >> Av1Constants.ModeInfoSizeLog2;
-                    if (childPosition.Y < this.picture.Parent.Common.ModeInfoRowCount &&
-                        childPosition.X < this.picture.Parent.Common.ModeInfoColumnCount)
-                    {
-                        this.SetBlockGeometry(childOrigin, Av1BlockSize.Block4x4, Av1PartitionType.None);
-                    }
-                }
+            bool hasColumns =
+                modeInfoPosition.X + blockSize.Get4x4WideCount() <= this.picture.Parent.Common.ModeInfoColumnCount;
 
-                return Av1PartitionType.Split;
+            if (!hasRows || !hasColumns)
+            {
+                // Coded dimensions are aligned to eight samples, so an incomplete searched node must retain
+                // the prepared split tree rather than evaluating a block that extends beyond source storage.
+                this.PreparePartitionGeometry(blockOrigin, blockSize, preparedPartition);
+                return preparedPartition;
             }
 
             if (this.effort < 9)
@@ -198,68 +208,102 @@ internal static partial class Av1IntraSuperblockEncoder
                 return preparedPartition;
             }
 
-            int savedLumaArea = this.codedAreaLuma;
-            int savedChromaArea = this.codedAreaChroma;
-            this.SavePartitionTrialContexts(blockOrigin, tileIndex);
-            long bestCost = this.EvaluatePartitionCandidate(
+            Av1PartitionType selectedPartition = this.SelectBestPartition(
                 writer,
                 macroBlock,
                 blockOrigin,
                 tileIndex,
-                blockSize,
-                Av1PartitionType.None);
+                blockSize);
 
-            Av1PartitionType selectedPartition = Av1PartitionType.None;
-            this.ResetPartitionTrial(blockOrigin, tileIndex, savedLumaArea, savedChromaArea);
-            long splitCost = this.EvaluatePartitionCandidate(
-                writer,
-                macroBlock,
-                blockOrigin,
-                tileIndex,
-                blockSize,
-                Av1PartitionType.Split);
+            // Trial reconstruction and mode entries need no copy-back. The selected branch is evaluated again
+            // in raster order, overwriting each trial-local value before a later selected leaf can consume it.
+            this.PreparePartitionGeometry(blockOrigin, blockSize, selectedPartition);
+            return selectedPartition;
+        }
 
-            if (splitCost < bestCost)
+        private Av1PartitionType SelectBestPartition(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Point blockOrigin,
+            ushort tileIndex,
+            Av1BlockSize blockSize)
+        {
+            int savedLumaArea = this.codedAreaLuma;
+            int savedChromaArea = this.codedAreaChroma;
+            this.SavePartitionTrialContexts(blockOrigin, tileIndex, blockSize);
+            long bestCost = long.MaxValue;
+            Av1PartitionType selectedPartition = Av1PartitionType.None;
+            ReadOnlySpan<Av1PartitionType> searchOrder = PartitionSearchOrder;
+            int candidateCount = blockSize == Av1BlockSize.Block8x8 ? 4 : searchOrder.Length;
+            for (int candidateIndex = 0; candidateIndex < candidateCount; candidateIndex++)
             {
-                bestCost = splitCost;
-                selectedPartition = Av1PartitionType.Split;
+                Av1PartitionType partitionType = searchOrder[candidateIndex];
+                if (!this.IsPartitionCandidateAllowed(blockSize, partitionType))
+                {
+                    continue;
+                }
+
+                long candidateCost = this.EvaluatePartitionCandidate(
+                    writer,
+                    macroBlock,
+                    blockOrigin,
+                    tileIndex,
+                    blockSize,
+                    partitionType,
+                    publishFinalContexts: false);
+
+                if (candidateCost < bestCost)
+                {
+                    bestCost = candidateCost;
+                    selectedPartition = partitionType;
+                }
+
+                this.ResetPartitionTrial(
+                    blockOrigin,
+                    tileIndex,
+                    blockSize,
+                    savedLumaArea,
+                    savedChromaArea);
             }
 
-            this.ResetPartitionTrial(blockOrigin, tileIndex, savedLumaArea, savedChromaArea);
-            long horizontalCost = this.EvaluatePartitionCandidate(
+            return selectedPartition;
+        }
+
+        private long EvaluateSelectedPartitionTree(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Point blockOrigin,
+            ushort tileIndex,
+            Av1BlockSize blockSize,
+            bool publishContexts)
+        {
+            Av1PartitionType selectedPartition = this.SelectBestPartition(
                 writer,
                 macroBlock,
                 blockOrigin,
                 tileIndex,
-                blockSize,
-                Av1PartitionType.Horizontal);
-
-            if (horizontalCost < bestCost)
-            {
-                bestCost = horizontalCost;
-                selectedPartition = Av1PartitionType.Horizontal;
-            }
+                blockSize);
 
-            this.ResetPartitionTrial(blockOrigin, tileIndex, savedLumaArea, savedChromaArea);
-            long verticalCost = this.EvaluatePartitionCandidate(
+            long cost = this.EvaluatePartitionCandidate(
                 writer,
                 macroBlock,
                 blockOrigin,
                 tileIndex,
                 blockSize,
-                Av1PartitionType.Vertical);
+                selectedPartition,
+                publishContexts);
 
-            if (verticalCost < bestCost)
+            if (publishContexts)
             {
-                selectedPartition = Av1PartitionType.Vertical;
+                Av1TileWriter.UpdatePartitionContexts(
+                    this.picture.PartitionContexts[tileIndex],
+                    blockOrigin,
+                    selectedPartition.GetBlockSubSize(blockSize),
+                    blockSize,
+                    selectedPartition);
             }
 
-            this.ResetPartitionTrial(blockOrigin, tileIndex, savedLumaArea, savedChromaArea);
-
-            // Trial reconstruction and mode entries need no copy-back. The selected branch is evaluated again
-            // in raster order, overwriting each trial-local value before a later selected leaf can consume it.
-            this.PreparePartitionGeometry(blockOrigin, blockSize, selectedPartition);
-            return selectedPartition;
+            return cost;
         }
 
         private long EvaluatePartitionCandidate(
@@ -268,7 +312,8 @@ internal static partial class Av1IntraSuperblockEncoder
             Point blockOrigin,
             ushort tileIndex,
             Av1BlockSize blockSize,
-            Av1PartitionType partitionType)
+            Av1PartitionType partitionType,
+            bool publishFinalContexts)
         {
             int rate = Av1TileWriter.GetPartitionCost(
                 this.picture,
@@ -279,31 +324,36 @@ internal static partial class Av1IntraSuperblockEncoder
                 this.picture.PartitionContexts[tileIndex]);
 
             long cost = Av1RateDistortion.GetCost(this.rateMultiplier, rate, 0);
-            int leafCount = partitionType == Av1PartitionType.Split
-                ? 4
-                : partitionType == Av1PartitionType.None
-                    ? 1
-                    : 2;
-
-            Av1BlockSize leafSize = partitionType.GetBlockSubSize(blockSize);
+            int leafCount = GetPartitionLeafCount(partitionType);
 
             // Child reconstruction and syntax contexts become input to the next child. Publishing only
-            // non-final leaves reproduces libaom's raster trial without writing entropy symbols.
+            // the required leaves reproduces libaom's raster dry run without writing entropy symbols.
             for (int leafIndex = 0; leafIndex < leafCount; leafIndex++)
             {
-                Point leafOrigin = GetPartitionLeafOrigin(
+                GetPartitionLeafGeometry(
                     blockOrigin,
                     blockSize,
                     partitionType,
-                    leafIndex);
+                    leafIndex,
+                    out Point leafOrigin,
+                    out Av1BlockSize leafSize);
 
-                cost += this.EvaluatePartitionLeaf(
-                    writer,
-                    macroBlock,
-                    leafOrigin,
-                    tileIndex,
-                    leafSize,
-                    publishContexts: leafIndex < leafCount - 1);
+                bool publishContexts = leafIndex < leafCount - 1 || publishFinalContexts;
+                cost += partitionType == Av1PartitionType.Split && blockSize > Av1BlockSize.Block8x8
+                    ? this.EvaluateSelectedPartitionTree(
+                        writer,
+                        macroBlock,
+                        leafOrigin,
+                        tileIndex,
+                        leafSize,
+                        publishContexts)
+                    : this.EvaluatePartitionLeaf(
+                        writer,
+                        macroBlock,
+                        leafOrigin,
+                        tileIndex,
+                        leafSize,
+                        publishContexts);
             }
 
             return cost;
@@ -312,12 +362,13 @@ internal static partial class Av1IntraSuperblockEncoder
         private void ResetPartitionTrial(
             Point blockOrigin,
             ushort tileIndex,
+            Av1BlockSize blockSize,
             int savedLumaArea,
             int savedChromaArea)
         {
             this.codedAreaLuma = savedLumaArea;
             this.codedAreaChroma = savedChromaArea;
-            this.RestorePartitionTrialContexts(blockOrigin, tileIndex);
+            this.RestorePartitionTrialContexts(blockOrigin, tileIndex, blockSize);
         }
 
         private void PreparePartitionGeometry(
@@ -325,42 +376,149 @@ internal static partial class Av1IntraSuperblockEncoder
             Av1BlockSize blockSize,
             Av1PartitionType partitionType)
         {
-            int leafCount = partitionType == Av1PartitionType.Split
-                ? 4
-                : partitionType == Av1PartitionType.None
-                    ? 1
-                    : 2;
-
-            Av1BlockSize leafSize = partitionType.GetBlockSubSize(blockSize);
+            int leafCount = GetPartitionLeafCount(partitionType);
             for (int leafIndex = 0; leafIndex < leafCount; leafIndex++)
             {
-                Point leafOrigin = GetPartitionLeafOrigin(
+                GetPartitionLeafGeometry(
                     blockOrigin,
                     blockSize,
                     partitionType,
-                    leafIndex);
+                    leafIndex,
+                    out Point leafOrigin,
+                    out Av1BlockSize leafSize);
 
-                this.SetBlockGeometry(leafOrigin, leafSize, Av1PartitionType.None);
+                if (this.IsBlockOriginInsideFrame(leafOrigin))
+                {
+                    this.SetBlockGeometry(leafOrigin, leafSize, Av1PartitionType.None);
+                }
             }
         }
 
-        private static Point GetPartitionLeafOrigin(
+        private bool IsPartitionCandidateAllowed(
+            Av1BlockSize blockSize,
+            Av1PartitionType partitionType)
+        {
+            if (partitionType.GetBlockSubSize(blockSize) == Av1BlockSize.Invalid)
+            {
+                return false;
+            }
+
+            if (this.source.IsMonochrome)
+            {
+                return true;
+            }
+
+            ObuColorConfig colorConfig = this.picture.Sequence.SequenceHeader.ColorConfig;
+            int leafCount = GetPartitionLeafCount(partitionType);
+            for (int leafIndex = 0; leafIndex < leafCount; leafIndex++)
+            {
+                GetPartitionLeafGeometry(
+                    Point.Empty,
+                    blockSize,
+                    partitionType,
+                    leafIndex,
+                    out _,
+                    out Av1BlockSize leafSize);
+
+                if (leafSize.GetSubsampled(colorConfig.SubSamplingX, colorConfig.SubSamplingY) ==
+                    Av1BlockSize.Invalid)
+                {
+                    return false;
+                }
+            }
+
+            return true;
+        }
+
+        private bool IsBlockOriginInsideFrame(Point blockOrigin)
+        {
+            Point modeInfoPosition = blockOrigin >> Av1Constants.ModeInfoSizeLog2;
+            return modeInfoPosition.Y < this.picture.Parent.Common.ModeInfoRowCount &&
+                modeInfoPosition.X < this.picture.Parent.Common.ModeInfoColumnCount;
+        }
+
+        private static int GetPartitionLeafCount(Av1PartitionType partitionType)
+            => partitionType switch
+            {
+                Av1PartitionType.None => 1,
+                Av1PartitionType.Horizontal or Av1PartitionType.Vertical => 2,
+                Av1PartitionType.HorizontalA or
+                    Av1PartitionType.HorizontalB or
+                    Av1PartitionType.VerticalA or
+                    Av1PartitionType.VerticalB => 3,
+                _ => 4
+            };
+
+        private static void GetPartitionLeafGeometry(
             Point blockOrigin,
             Av1BlockSize blockSize,
             Av1PartitionType partitionType,
-            int leafIndex)
+            int leafIndex,
+            out Point leafOrigin,
+            out Av1BlockSize leafSize)
         {
             int halfWidth = blockSize.GetWidth() >> 1;
             int halfHeight = blockSize.GetHeight() >> 1;
-            return partitionType switch
+            Av1BlockSize rectangularSize = partitionType.GetBlockSubSize(blockSize);
+            Av1BlockSize splitSize = Av1PartitionType.Split.GetBlockSubSize(blockSize);
+            switch (partitionType)
             {
-                Av1PartitionType.Horizontal => blockOrigin + new Size(0, leafIndex * halfHeight),
-                Av1PartitionType.Vertical => blockOrigin + new Size(leafIndex * halfWidth, 0),
-                Av1PartitionType.Split => blockOrigin + new Size(
-                    (leafIndex & 1) * halfWidth,
-                    (leafIndex >> 1) * halfHeight),
-                _ => blockOrigin
-            };
+                case Av1PartitionType.Horizontal:
+                    leafOrigin = blockOrigin + new Size(0, leafIndex * halfHeight);
+                    leafSize = rectangularSize;
+                    return;
+                case Av1PartitionType.Vertical:
+                    leafOrigin = blockOrigin + new Size(leafIndex * halfWidth, 0);
+                    leafSize = rectangularSize;
+                    return;
+                case Av1PartitionType.Split:
+                    leafOrigin = blockOrigin + new Size(
+                        (leafIndex & 1) * halfWidth,
+                        (leafIndex >> 1) * halfHeight);
+
+                    leafSize = splitSize;
+                    return;
+                case Av1PartitionType.HorizontalA:
+                    leafOrigin = leafIndex < 2
+                        ? blockOrigin + new Size(leafIndex * halfWidth, 0)
+                        : blockOrigin + new Size(0, halfHeight);
+
+                    leafSize = leafIndex < 2 ? splitSize : rectangularSize;
+                    return;
+                case Av1PartitionType.HorizontalB:
+                    leafOrigin = leafIndex == 0
+                        ? blockOrigin
+                        : blockOrigin + new Size((leafIndex - 1) * halfWidth, halfHeight);
+
+                    leafSize = leafIndex == 0 ? rectangularSize : splitSize;
+                    return;
+                case Av1PartitionType.VerticalA:
+                    leafOrigin = leafIndex < 2
+                        ? blockOrigin + new Size(0, leafIndex * halfHeight)
+                        : blockOrigin + new Size(halfWidth, 0);
+
+                    leafSize = leafIndex < 2 ? splitSize : rectangularSize;
+                    return;
+                case Av1PartitionType.VerticalB:
+                    leafOrigin = leafIndex == 0
+                        ? blockOrigin
+                        : blockOrigin + new Size(halfWidth, (leafIndex - 1) * halfHeight);
+
+                    leafSize = leafIndex == 0 ? rectangularSize : splitSize;
+                    return;
+                case Av1PartitionType.Horizontal4:
+                    leafOrigin = blockOrigin + new Size(0, leafIndex * (blockSize.GetHeight() >> 2));
+                    leafSize = rectangularSize;
+                    return;
+                case Av1PartitionType.Vertical4:
+                    leafOrigin = blockOrigin + new Size(leafIndex * (blockSize.GetWidth() >> 2), 0);
+                    leafSize = rectangularSize;
+                    return;
+                default:
+                    leafOrigin = blockOrigin;
+                    leafSize = blockSize;
+                    return;
+            }
         }
 
         /// <inheritdoc/>
@@ -658,7 +816,8 @@ internal static partial class Av1IntraSuperblockEncoder
                     lumaArea,
                     chromaArea,
                     modeInfo,
-                    block);
+                    block,
+                    paletteInfo);
             }
 
             return this.selectedBlockCost;
@@ -686,7 +845,8 @@ internal static partial class Av1IntraSuperblockEncoder
             int lumaArea,
             int chromaArea,
             Av1MacroBlockModeInfo modeInfo,
-            Av1EncoderBlockStruct block)
+            Av1EncoderBlockStruct block,
+            Av1EncoderPaletteInfo paletteInfo)
         {
             Av1BlockSize blockSize = modeInfo.Block.BlockSize;
             Av1TransformSize transformSize = modeInfo.Block.TransformSize;
@@ -707,21 +867,27 @@ internal static partial class Av1IntraSuperblockEncoder
             Span<Av1EncoderTransformBlockState> lumaStates =
                 this.coefficientBuffer.GetTransformBlockSpan(this.superblock.Index, Av1Plane.Y);
 
-            Av1EncoderTransformBlockState lumaState =
-                lumaStates[lumaArea / Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount];
-
             Span<int> lumaCoefficients = this.coefficientBuffer.GetPlaneSpan(this.superblock.Index, Av1Plane.Y);
-            byte lumaContext = Av1SymbolContextHelper.GetCoefficientContext(
-                lumaCoefficients[lumaArea..],
+            PublishCoefficientContexts(
+                this.picture.LuminanceDcSignLevelCoefficientNeighbors[tileIndex],
+                blockOrigin,
+                blockSize,
                 transformSize,
-                lumaState.TransformType,
-                lumaState.EndOfBlock);
+                lumaCoefficients[lumaArea..],
+                lumaStates[(lumaArea / Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount)..]);
 
-            this.picture.LuminanceDcSignLevelCoefficientNeighbors[tileIndex].UnitModeWrite(
-                lumaContext,
-                blockOrigin,
-                blockDimensions,
-                Av1NeighborArrayUnit<byte>.UnitMask.Top | Av1NeighborArrayUnit<byte>.UnitMask.Left);
+            if (this.picture.Parent.FrameHeader.AllowScreenContentTools)
+            {
+                const Av1NeighborArrayUnit<Av1EncoderPaletteInfo>.UnitMask PaletteContextMask =
+                    Av1NeighborArrayUnit<Av1EncoderPaletteInfo>.UnitMask.Top |
+                    Av1NeighborArrayUnit<Av1EncoderPaletteInfo>.UnitMask.Left;
+
+                this.picture.PaletteContexts[tileIndex].UnitModeWrite(
+                    paletteInfo,
+                    blockOrigin,
+                    blockDimensions,
+                    PaletteContextMask);
+            }
 
             if (!block.HasChroma)
             {
@@ -753,56 +919,111 @@ internal static partial class Av1IntraSuperblockEncoder
             Span<Av1EncoderTransformBlockState> redStates =
                 this.coefficientBuffer.GetTransformBlockSpan(this.superblock.Index, Av1Plane.V);
 
-            Av1EncoderTransformBlockState blueState = blueStates[chromaStateIndex];
-            Av1EncoderTransformBlockState redState = redStates[chromaStateIndex];
             Span<int> blueCoefficients = this.coefficientBuffer.GetPlaneSpan(this.superblock.Index, Av1Plane.U);
             Span<int> redCoefficients = this.coefficientBuffer.GetPlaneSpan(this.superblock.Index, Av1Plane.V);
-            byte blueContext = Av1SymbolContextHelper.GetCoefficientContext(
-                blueCoefficients[chromaArea..],
+            PublishCoefficientContexts(
+                this.picture.CbDcSignLevelCoefficientNeighbors[tileIndex],
+                chromaOrigin,
+                chromaBlockSize,
                 chromaTransformSize,
-                blueState.TransformType,
-                blueState.EndOfBlock);
+                blueCoefficients[chromaArea..],
+                blueStates[chromaStateIndex..]);
 
-            byte redContext = Av1SymbolContextHelper.GetCoefficientContext(
-                redCoefficients[chromaArea..],
+            PublishCoefficientContexts(
+                this.picture.CrDcSignLevelCoefficientNeighbors[tileIndex],
+                chromaOrigin,
+                chromaBlockSize,
                 chromaTransformSize,
-                redState.TransformType,
-                redState.EndOfBlock);
+                redCoefficients[chromaArea..],
+                redStates[chromaStateIndex..]);
+        }
 
-            Size chromaDimensions = new(chromaBlockSize.GetWidth(), chromaBlockSize.GetHeight());
-            this.picture.CbDcSignLevelCoefficientNeighbors[tileIndex].UnitModeWrite(
-                blueContext,
-                chromaOrigin,
-                chromaDimensions,
-                Av1NeighborArrayUnit<byte>.UnitMask.Top | Av1NeighborArrayUnit<byte>.UnitMask.Left);
+        private static void PublishCoefficientContexts(
+            Av1NeighborArrayUnit<byte> neighbors,
+            Point blockOrigin,
+            Av1BlockSize blockSize,
+            Av1TransformSize transformSize,
+            ReadOnlySpan<int> coefficients,
+            ReadOnlySpan<Av1EncoderTransformBlockState> states)
+        {
+            const Av1NeighborArrayUnit<byte>.UnitMask EdgeMask =
+                Av1NeighborArrayUnit<byte>.UnitMask.Top |
+                Av1NeighborArrayUnit<byte>.UnitMask.Left;
 
-            this.picture.CrDcSignLevelCoefficientNeighbors[tileIndex].UnitModeWrite(
-                redContext,
-                chromaOrigin,
-                chromaDimensions,
-                Av1NeighborArrayUnit<byte>.UnitMask.Top | Av1NeighborArrayUnit<byte>.UnitMask.Left);
+            int blockWidth = blockSize.GetWidth();
+            int blockHeight = blockSize.GetHeight();
+            int transformWidth = transformSize.GetWidth();
+            int transformHeight = transformSize.GetHeight();
+            int transformSampleCount = transformSize.GetSize2d();
+            int transformIndex = 0;
+            int coefficientOffset = 0;
+
+            // Uniform transform blocks are retained and written in raster order. Publishing that same tiling
+            // preserves the distinct top and left contexts consumed by the next coding block in a dry run.
+            for (int row = 0; row < blockHeight; row += transformHeight)
+            {
+                for (int column = 0; column < blockWidth; column += transformWidth)
+                {
+                    Av1EncoderTransformBlockState state = states[transformIndex++];
+                    byte context = Av1SymbolContextHelper.GetCoefficientContext(
+                        coefficients[coefficientOffset..],
+                        transformSize,
+                        state.TransformType,
+                        state.EndOfBlock);
+
+                    neighbors.UnitModeWrite(
+                        context,
+                        blockOrigin + new Size(column, row),
+                        new Size(transformWidth, transformHeight),
+                        EdgeMask);
+
+                    coefficientOffset += transformSampleCount;
+                }
+            }
         }
 
-        private void SavePartitionTrialContexts(Point blockOrigin, ushort tileIndex)
+        private void SavePartitionTrialContexts(
+            Point blockOrigin,
+            ushort tileIndex,
+            Av1BlockSize blockSize)
         {
-            Span<byte> storage = this.blockWorkspace.PartitionContexts;
+            Span<byte> storage = this.blockWorkspace.GetPartitionContextStorage(blockSize);
             int offset = 0;
+            SaveNeighborEdges(
+                this.picture.PartitionContexts[tileIndex],
+                blockOrigin,
+                blockSize.Get4x4WideCount(),
+                blockSize.Get4x4HighCount(),
+                storage,
+                ref offset);
+
             SaveNeighborEdges(
                 this.picture.LuminanceDcSignLevelCoefficientNeighbors[tileIndex],
                 blockOrigin,
-                Av1BlockSize.Block8x8.Get4x4WideCount(),
-                Av1BlockSize.Block8x8.Get4x4HighCount(),
+                blockSize.Get4x4WideCount(),
+                blockSize.Get4x4HighCount(),
                 storage,
                 ref offset);
 
             SaveNeighborEdges(
                 this.picture.TransformFunctionContexts[tileIndex],
                 blockOrigin,
-                Av1BlockSize.Block8x8.Get4x4WideCount(),
-                Av1BlockSize.Block8x8.Get4x4HighCount(),
+                blockSize.Get4x4WideCount(),
+                blockSize.Get4x4HighCount(),
                 storage,
                 ref offset);
 
+            if (this.picture.Parent.FrameHeader.AllowScreenContentTools)
+            {
+                SaveNeighborEdges(
+                    this.picture.PaletteContexts[tileIndex],
+                    blockOrigin,
+                    blockSize.Get4x4WideCount(),
+                    blockSize.Get4x4HighCount(),
+                    storage,
+                    ref offset);
+            }
+
             if (this.source.IsMonochrome)
             {
                 return;
@@ -816,7 +1037,7 @@ internal static partial class Av1IntraSuperblockEncoder
                 subsamplingX,
                 subsamplingY);
 
-            Av1BlockSize chromaBlockSize = Av1BlockSize.Block8x8.GetSubsampled(
+            Av1BlockSize chromaBlockSize = blockSize.GetSubsampled(
                 colorConfig.SubSamplingX,
                 colorConfig.SubSamplingY);
 
@@ -837,26 +1058,48 @@ internal static partial class Av1IntraSuperblockEncoder
                 ref offset);
         }
 
-        private void RestorePartitionTrialContexts(Point blockOrigin, ushort tileIndex)
+        private void RestorePartitionTrialContexts(
+            Point blockOrigin,
+            ushort tileIndex,
+            Av1BlockSize blockSize)
         {
-            ReadOnlySpan<byte> storage = this.blockWorkspace.PartitionContexts;
+            ReadOnlySpan<byte> storage = this.blockWorkspace.GetPartitionContextStorage(blockSize);
             int offset = 0;
+            RestoreNeighborEdges(
+                this.picture.PartitionContexts[tileIndex],
+                blockOrigin,
+                blockSize.Get4x4WideCount(),
+                blockSize.Get4x4HighCount(),
+                storage,
+                ref offset);
+
             RestoreNeighborEdges(
                 this.picture.LuminanceDcSignLevelCoefficientNeighbors[tileIndex],
                 blockOrigin,
-                Av1BlockSize.Block8x8.Get4x4WideCount(),
-                Av1BlockSize.Block8x8.Get4x4HighCount(),
+                blockSize.Get4x4WideCount(),
+                blockSize.Get4x4HighCount(),
                 storage,
                 ref offset);
 
             RestoreNeighborEdges(
                 this.picture.TransformFunctionContexts[tileIndex],
                 blockOrigin,
-                Av1BlockSize.Block8x8.Get4x4WideCount(),
-                Av1BlockSize.Block8x8.Get4x4HighCount(),
+                blockSize.Get4x4WideCount(),
+                blockSize.Get4x4HighCount(),
                 storage,
                 ref offset);
 
+            if (this.picture.Parent.FrameHeader.AllowScreenContentTools)
+            {
+                RestoreNeighborEdges(
+                    this.picture.PaletteContexts[tileIndex],
+                    blockOrigin,
+                    blockSize.Get4x4WideCount(),
+                    blockSize.Get4x4HighCount(),
+                    storage,
+                    ref offset);
+            }
+
             if (this.source.IsMonochrome)
             {
                 return;
@@ -870,7 +1113,7 @@ internal static partial class Av1IntraSuperblockEncoder
                 subsamplingX,
                 subsamplingY);
 
-            Av1BlockSize chromaBlockSize = Av1BlockSize.Block8x8.GetSubsampled(
+            Av1BlockSize chromaBlockSize = blockSize.GetSubsampled(
                 colorConfig.SubSamplingX,
                 colorConfig.SubSamplingY);
 
@@ -891,32 +1134,46 @@ internal static partial class Av1IntraSuperblockEncoder
                 ref offset);
         }
 
-        private static void SaveNeighborEdges(
-            Av1NeighborArrayUnit<byte> neighbors,
+        private static void SaveNeighborEdges<T>(
+            Av1NeighborArrayUnit<T> neighbors,
             Point blockOrigin,
             int width,
             int height,
             Span<byte> storage,
             ref int offset)
+            where T : struct
         {
-            neighbors.Top.Slice(neighbors.GetTopIndex(blockOrigin), width).CopyTo(storage[offset..]);
-            offset += width;
-            neighbors.Left.Slice(neighbors.GetLeftIndex(blockOrigin), height).CopyTo(storage[offset..]);
-            offset += height;
+            Span<byte> top = MemoryMarshal.AsBytes(
+                neighbors.Top.Slice(neighbors.GetTopIndex(blockOrigin), width));
+
+            top.CopyTo(storage[offset..]);
+            offset += top.Length;
+            Span<byte> left = MemoryMarshal.AsBytes(
+                neighbors.Left.Slice(neighbors.GetLeftIndex(blockOrigin), height));
+
+            left.CopyTo(storage[offset..]);
+            offset += left.Length;
         }
 
-        private static void RestoreNeighborEdges(
-            Av1NeighborArrayUnit<byte> neighbors,
+        private static void RestoreNeighborEdges<T>(
+            Av1NeighborArrayUnit<T> neighbors,
             Point blockOrigin,
             int width,
             int height,
             ReadOnlySpan<byte> storage,
             ref int offset)
+            where T : struct
         {
-            storage.Slice(offset, width).CopyTo(neighbors.Top[neighbors.GetTopIndex(blockOrigin)..]);
-            offset += width;
-            storage.Slice(offset, height).CopyTo(neighbors.Left[neighbors.GetLeftIndex(blockOrigin)..]);
-            offset += height;
+            Span<byte> top = MemoryMarshal.AsBytes(
+                neighbors.Top.Slice(neighbors.GetTopIndex(blockOrigin), width));
+
+            storage.Slice(offset, top.Length).CopyTo(top);
+            offset += top.Length;
+            Span<byte> left = MemoryMarshal.AsBytes(
+                neighbors.Left.Slice(neighbors.GetLeftIndex(blockOrigin), height));
+
+            storage.Slice(offset, left.Length).CopyTo(left);
+            offset += left.Length;
         }
 
         private long GetRegularBlockCost(

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PartitionContext.cs

@@ -2,6 +2,7 @@
 // Licensed under the Six Labors Split License.
 
 using System.Numerics;
+using System.Runtime.InteropServices;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
@@ -12,8 +13,14 @@ namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 /// Each set bit records a split at one block-size level. For example, <c>11111</c> records splits from
 /// 128 by 128 through 8 by 8, while <c>10000</c> records only the 128 by 128 split.
 /// </remarks>
+[StructLayout(LayoutKind.Sequential, Size = StorageSize)]
 internal struct Av1PartitionContext : IMinMaxValue<Av1PartitionContext>
 {
+    /// <summary>
+    /// The packed size of the above and left context bytes.
+    /// </summary>
+    public const int StorageSize = 2;
+
     /// <summary>
     /// Maps each block size to the five-bit context stored for an above neighbor.
     /// </summary>

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

@@ -216,6 +216,64 @@ public class Av1EncoderFrameTests
         Assert.Equal(new Size(Size, Size), decoded.Size);
     }
 
+    [Fact]
+    public void EncodeEffortNineSelectsSixteenBySixteenVerticalPartition()
+    {
+        const int Size = 32;
+        using Image<Rgba32> source = new(Size, Size);
+        for (int y = 0; y < Size; y++)
+        {
+            Span<Rgba32> row = source.Frames.RootFrame.PixelBuffer.DangerousGetRowSpan(y);
+            for (int x = 0; x < Size; x++)
+            {
+                byte value = 128;
+                if (x == 15 && y >= 16)
+                {
+                    value = (byte)(24 + ((y - 16) * 13));
+                }
+                else if (y == 15 && x >= 16)
+                {
+                    value = (byte)(16 + ((x - 16) * 15));
+                }
+                else if (x >= 16 && y >= 16)
+                {
+                    // The left 8x16 half repeats its external left edge, while the right half repeats
+                    // its external top edge. One 16x16 predictor cannot reproduce both surfaces.
+                    value = x < 24
+                        ? (byte)(24 + ((y - 16) * 13))
+                        : (byte)(16 + ((x - 16) * 15));
+                }
+
+                row[x] = new Rgba32(value, value, value);
+            }
+        }
+
+        using MemoryStream stream = new();
+        _ = Av1FrameEncoder.Encode(
+            Configuration.Default,
+            source.Frames.RootFrame,
+            stream,
+            CreateColorConfig(Av1BitDepth.EightBit, Av1ColorFormat.Yuv400),
+            qIndex: 4,
+            effort: 9);
+
+        byte[] payload = stream.ToArray();
+        using Av1Decoder decoder = new(Configuration.Default);
+        using Image<Rgba32> decoded = decoder.Decode<Rgba32>(payload);
+        Av1FrameInfo frameInfo = Assert.IsType<Av1FrameInfo>(decoder.FrameInfo);
+        for (int modeInfoY = 4; modeInfoY < 8; modeInfoY++)
+        {
+            for (int modeInfoX = 4; modeInfoX < 8; modeInfoX++)
+            {
+                Assert.Equal(
+                    Av1BlockSize.Block8x16,
+                    frameInfo.GetModeInfoAt(new Point(modeInfoX, modeInfoY)).BlockSize);
+            }
+        }
+
+        Assert.Equal(new Size(Size, Size), decoded.Size);
+    }
+
     [Theory]
     [InlineData(EightBit)]
     [InlineData(TenBit)]

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

@@ -617,10 +617,17 @@ public class Av1TransformBlockEncoderTests
             Av1EncoderIntraBlockCopyWorkspace<ushort> intraBlockCopyWorkspace =
                 workspace.GetIntraBlockCopyWorkspace<ushort>();
 
-            Assert.Equal(17, modeWorkspace.GetReferenceSamples(3).Length);
+            Assert.Equal(
+                (2 * Av1EncoderModeDecisionWorkspace<ushort>.MaximumBlockDimension) + 1,
+                modeWorkspace.GetReferenceSamples(3).Length);
+
             Assert.Equal(Av1EncoderModeDecisionWorkspace<ushort>.MaximumSampleCount, modeWorkspace.GetCandidateReconstruction(1).Length);
             Assert.Equal(Av1EncoderModeDecisionWorkspace<ushort>.MaximumSampleCount, modeWorkspace.GetCandidateCoefficients(1).Length);
-            Assert.Equal(Av1ChromaFromLumaContext.BufferLine * 8, modeWorkspace.ChromaFromLumaSamples.Length);
+            Assert.Equal(
+                Av1ChromaFromLumaContext.BufferLine *
+                    Av1EncoderModeDecisionWorkspace<ushort>.MaximumBlockDimension,
+                modeWorkspace.ChromaFromLumaSamples.Length);
+
             Assert.Equal(Av1ChromaFromLumaMath.AlphaCandidateCount, modeWorkspace.GetChromaFromLumaRates(1).Length);
             Assert.Equal(Av1ChromaFromLumaMath.AlphaCandidateCount, modeWorkspace.GetChromaFromLumaDistortions(1).Length);
             Assert.Equal(Av1EncoderModeDecisionWorkspace<ushort>.MaximumSampleCount, paletteWorkspace.GetPrediction(1).Length);