Implement live AV1 8x8 partition search

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

@@ -722,6 +722,15 @@ internal class Av1SymbolEncoder : IDisposable
     public int GetDisplacementVectorSearchCost(Av1MotionVector value, Av1MotionVector reference)
         => this.displacementVector.GetCost(this.writer, value, reference);
 
+    /// <summary>
+    /// Gets the current fixed-point cost of a complete block partition symbol.
+    /// </summary>
+    /// <param name="partitionType">The partition type to measure.</param>
+    /// <param name="context">The partition probability context.</param>
+    /// <returns>The rate cost in 1/512-bit units.</returns>
+    public int GetPartitionTypeCost(Av1PartitionType partitionType, int context)
+        => Av1ProbabilityCost.GetSymbolCost(this.tilePartitionTypes[context], (int)partitionType);
+
     /// <summary>
     /// Writes a complete block partition type using the selected partition context.
     /// </summary>
@@ -747,6 +756,26 @@ internal class Av1SymbolEncoder : IDisposable
         w.WriteBoolean(value, frequency);
     }
 
+    /// <summary>
+    /// Gets the current fixed-point cost of the split-versus-horizontal boundary decision.
+    /// </summary>
+    /// <param name="partitionType">The split or horizontal partition outcome.</param>
+    /// <param name="blockSize">The current block size.</param>
+    /// <param name="context">The partition probability context.</param>
+    /// <returns>The rate cost in 1/512-bit units.</returns>
+    public int GetSplitOrHorizontalCost(Av1PartitionType partitionType, Av1BlockSize blockSize, int context)
+    {
+        int frequency = (int)Av1SymbolDecoder.GetSplitOrHorizontalFrequency(
+            this.tilePartitionTypes,
+            blockSize,
+            context);
+
+        return Av1ProbabilityCost.GetProbabilityCost(
+            partitionType == Av1PartitionType.Split
+                ? frequency
+                : Av1Distribution.ProbabilityTop - frequency);
+    }
+
     /// <summary>
     /// Writes the split-versus-vertical boundary decision for a block clipped at the right tile edge.
     /// </summary>
@@ -761,6 +790,26 @@ internal class Av1SymbolEncoder : IDisposable
         w.WriteBoolean(value, frequency);
     }
 
+    /// <summary>
+    /// Gets the current fixed-point cost of the split-versus-vertical boundary decision.
+    /// </summary>
+    /// <param name="partitionType">The split or vertical partition outcome.</param>
+    /// <param name="blockSize">The current block size.</param>
+    /// <param name="context">The partition probability context.</param>
+    /// <returns>The rate cost in 1/512-bit units.</returns>
+    public int GetSplitOrVerticalCost(Av1PartitionType partitionType, Av1BlockSize blockSize, int context)
+    {
+        int frequency = (int)Av1SymbolDecoder.GetSplitOrVerticalFrequency(
+            this.tilePartitionTypes,
+            blockSize,
+            context);
+
+        return Av1ProbabilityCost.GetProbabilityCost(
+            partitionType == Av1PartitionType.Split
+                ? frequency
+                : Av1Distribution.ProbabilityTop - frequency);
+    }
+
     /// <summary>
     /// Encodes one transform block's coefficient syntax using scan-order probability contexts.
     /// </summary>

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

@@ -31,7 +31,8 @@ internal sealed class Av1EncoderBlockWorkspace : IDisposable
         MaximumCoefficientCount +
         MaximumCoefficientCount +
         Av1TransformWorkspace.MaximumLength +
-        IntraBlockCopyStorageLength;
+        IntraBlockCopyStorageLength +
+        PartitionContextStorageLength;
 
     private const int ResidualStorageLength = MaximumResidualCount / 2;
     private const int TransformCoefficientOffset = ResidualStorageLength;
@@ -64,6 +65,11 @@ internal sealed class Av1EncoderBlockWorkspace : IDisposable
         IntraBlockCopyResidualStorageLength +
         IntraBlockCopyCoefficientStorageLength;
 
+    private const int PartitionContextStorageOffset =
+        IntraBlockCopySampleStorageOffset + IntraBlockCopyStorageLength;
+
+    private const int PartitionContextStorageLength = 4;
+
     /// <summary>
     /// Owns the complete reusable block workspace in 32-bit elements so every transform region is naturally aligned.
     /// </summary>
@@ -100,6 +106,15 @@ internal sealed class Av1EncoderBlockWorkspace : IDisposable
     public Span<int> TransformWorkspace
         => this.owner.Memory.Span.Slice(TransformWorkspaceOffset, Av1TransformWorkspace.MaximumLength);
 
+    /// <summary>
+    /// Gets storage for the coefficient and transform edges restored after an 8x8 partition trial.
+    /// </summary>
+    public Span<byte> PartitionContexts
+        => MemoryMarshal.AsBytes(
+            this.owner.Memory.Span.Slice(
+                PartitionContextStorageOffset,
+                PartitionContextStorageLength));
+
     /// <summary>
     /// Gets the reusable storage used while comparing spatial, chroma-from-luma, filter-intra, and palette candidates.
     /// </summary>

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

@@ -300,7 +300,8 @@ internal static class Av1FrameEncoder
             sequenceHeader,
             frameHeader,
             image.Width,
-            image.Height);
+            image.Height,
+            disallow4x4AllFrames: effort < 9);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             configuration,
@@ -363,7 +364,8 @@ internal static class Av1FrameEncoder
             sequenceHeader,
             frameHeader,
             image.Width,
-            image.Height);
+            image.Height,
+            disallow4x4AllFrames: effort < 9);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             configuration,

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

@@ -46,6 +46,7 @@ internal static partial class Av1IntraSuperblockEncoder
             Av1MacroBlockModeInfo modeInfo,
             Point lumaOrigin,
             Point chromaOrigin,
+            Av1BlockSize blockSize,
             ushort tileIndex,
             Av1PredictionMode lumaMode,
             Av1TransformSize transformSize,
@@ -59,7 +60,6 @@ internal static partial class Av1IntraSuperblockEncoder
             out sbyte selectedChromaFromLumaSigns,
             out long selectedCost)
         {
-            const Av1BlockSize BlockSize = Av1BlockSize.Block8x8;
             Av1EncoderModeDecisionWorkspace<TSample> workspace =
                 this.blockWorkspace.GetModeDecisionWorkspace<TSample>();
 
@@ -73,11 +73,24 @@ internal static partial class Av1IntraSuperblockEncoder
             int modeInfoColumn = lumaOrigin.X >> Av1Constants.ModeInfoSizeLog2;
             bool hasLeft = macroBlock.IsLeftAvailable;
             bool hasAbove = macroBlock.IsUpAvailable;
+
+            // Subsampled chroma belongs to the bottom-right luma unit in its shared 8x8 region. Its external
+            // references therefore begin before that region, not immediately beside the owning 4x4 luma block.
+            if (subsamplingX != 0 && blockSize.Get4x4WideCount() < Av1BlockSize.Block8x8.Get4x4WideCount())
+            {
+                hasLeft = modeInfoColumn - 1 > macroBlock.Tile.ModeInfoColumnStart;
+            }
+
+            if (subsamplingY != 0 && blockSize.Get4x4HighCount() < Av1BlockSize.Block8x8.Get4x4HighCount())
+            {
+                hasAbove = modeInfoRow - 1 > macroBlock.Tile.ModeInfoRowStart;
+            }
+
             bool rightAvailable = modeInfoColumn + (transformSize.Get4x4WideCount() << subsamplingX) < macroBlock.Tile.ModeInfoColumnEnd;
             bool bottomAvailable = modeInfoRow + (transformSize.Get4x4HighCount() << subsamplingY) < macroBlock.Tile.ModeInfoRowEnd;
             bool hasTopRight = Av1IntraReferenceAvailability.HasTopRight(
                 this.picture.Sequence.SequenceHeader.SuperblockSize,
-                BlockSize,
+                blockSize,
                 modeInfoRow,
                 modeInfoColumn,
                 hasAbove,
@@ -91,7 +104,7 @@ internal static partial class Av1IntraSuperblockEncoder
 
             bool hasBottomLeft = Av1IntraReferenceAvailability.HasBottomLeft(
                 this.picture.Sequence.SequenceHeader.SuperblockSize,
-                BlockSize,
+                blockSize,
                 modeInfoRow,
                 modeInfoColumn,
                 bottomAvailable,
@@ -139,7 +152,7 @@ internal static partial class Av1IntraSuperblockEncoder
             ReadOnlySpan<TSample> blueLeft = blueLeftStorage.Slice(1, height * 2);
             ReadOnlySpan<TSample> redAbove = redAboveStorage.Slice(1, width * 2);
             ReadOnlySpan<TSample> redLeft = redLeftStorage.Slice(1, height * 2);
-            Av1BlockSize chromaBlockSize = BlockSize.GetSubsampled(colorConfig.SubSamplingX, colorConfig.SubSamplingY);
+            Av1BlockSize chromaBlockSize = blockSize.GetSubsampled(colorConfig.SubSamplingX, colorConfig.SubSamplingY);
             Av1TransformBlockContext blueContext = Av1TileWriter.GetTransformBlockContexts(
                 Av1ComponentType.Chroma,
                 this.picture.CbDcSignLevelCoefficientNeighbors[tileIndex],
@@ -172,11 +185,13 @@ internal static partial class Av1IntraSuperblockEncoder
             {
                 0 => 1,
                 1 => baseModeCount,
-                _ => baseModeCount + (directionalModeCount * deltaCount)
+                _ when blockSize >= Av1BlockSize.Block8x8 => baseModeCount + (directionalModeCount * deltaCount),
+                _ => baseModeCount
             };
 
             bool hasLumaPalette = paletteInfo.PaletteSizes[0] != 0;
-            int paletteDisabledCost = this.picture.Parent.FrameHeader.AllowScreenContentTools
+            int paletteDisabledCost = blockSize >= Av1BlockSize.Block8x8 &&
+                this.picture.Parent.FrameHeader.AllowScreenContentTools
                 ? writer.GetPaletteUvModeCost(false, hasLumaPalette)
                 : 0;
 
@@ -209,6 +224,7 @@ internal static partial class Av1IntraSuperblockEncoder
                     lumaMode,
                     chromaMode,
                     angleDelta,
+                    blockSize,
                     chromaOrigin,
                     transformSize,
                     blueSource,
@@ -257,7 +273,7 @@ internal static partial class Av1IntraSuperblockEncoder
                 }
             }
 
-            bool chromaFromLumaAllowed = BlockSize.AllowsChromaFromLuma(
+            bool chromaFromLumaAllowed = blockSize.AllowsChromaFromLuma(
                 this.picture.Parent.FrameHeader.LosslessArray[modeInfo.Block.SegmentId],
                 colorConfig.SubSamplingX,
                 colorConfig.SubSamplingY);
@@ -265,9 +281,15 @@ internal static partial class Av1IntraSuperblockEncoder
             if (this.effort >= 4 && chromaFromLumaAllowed)
             {
                 Span<short> lumaQ3 = workspace.ChromaFromLumaSamples;
+                Point chromaLumaOrigin = new(
+                    chromaOrigin.X << subsamplingX,
+                    chromaOrigin.Y << subsamplingY);
+
+                // CfL consumes the complete luma region represented by this chroma block, which starts before
+                // the bottom-right ownership point for subsampled 4x4 luma leaves.
                 TOperator.PrepareChromaFromLuma(
                     this.reconstruction.GetPlane(Av1Plane.Y),
-                    lumaOrigin,
+                    chromaLumaOrigin,
                     lumaQ3,
                     transformSize,
                     colorConfig.SubSamplingX,
@@ -345,7 +367,7 @@ internal static partial class Av1IntraSuperblockEncoder
                     this.picture.Parent.FrameHeader,
                     colorConfig,
                     modeInfo,
-                    BlockSize,
+                    blockSize,
                     lumaMode,
                     Av1ChromaPredictionMode.ChromaFromLuma,
                     0);
@@ -457,6 +479,7 @@ internal static partial class Av1IntraSuperblockEncoder
             }
 
             if (this.effort >= 5 &&
+                blockSize == Av1BlockSize.Block8x8 &&
                 this.picture.Parent.FrameHeader.AllowScreenContentTools &&
                 this.SelectChromaPalette(
                     writer,
@@ -544,6 +567,7 @@ internal static partial class Av1IntraSuperblockEncoder
             Av1PredictionMode lumaMode,
             Av1ChromaPredictionMode chromaMode,
             int angleDelta,
+            Av1BlockSize blockSize,
             Point chromaOrigin,
             Av1TransformSize transformSize,
             Buffer2DRegion<TSample> blueSource,
@@ -564,7 +588,6 @@ internal static partial class Av1IntraSuperblockEncoder
             ref Av1EncoderTransformBlockState candidateBlueState,
             ref Av1EncoderTransformBlockState candidateRedState)
         {
-            const Av1BlockSize BlockSize = Av1BlockSize.Block8x8;
             Av1PredictionMode predictionMode = chromaMode.ToLumaMode();
 
             // Intra chroma derives one transform type from the shared UV prediction mode. The type is not
@@ -623,7 +646,7 @@ internal static partial class Av1IntraSuperblockEncoder
                 this.picture.Parent.FrameHeader,
                 this.picture.Sequence.SequenceHeader.ColorConfig,
                 modeInfo,
-                BlockSize,
+                blockSize,
                 lumaMode,
                 chromaMode,
                 angleDelta);

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

@@ -20,13 +20,12 @@ internal static partial class Av1IntraSuperblockEncoder
         where TSample : unmanaged
         where TOperator : struct, IBlockEncodingOperator<TSample>
     {
-        private void SelectIntraBlockCopy(
+        private long SelectIntraBlockCopy(
             Av1SymbolEncoder writer,
             Av1MacroBlockD macroBlock,
             Point blockOrigin,
             ushort tileIndex,
-            long regularModeCost,
-            int regularEmptyTransformRate,
+            long regularCost,
             ref Av1MacroBlockModeInfo modeInfo,
             ref Av1EncoderBlockStruct block,
             ref Av1EncoderPaletteInfo paletteInfo)
@@ -99,21 +98,11 @@ internal static partial class Av1IntraSuperblockEncoder
 
             if (uniqueCandidateCount == 0)
             {
-                return;
+                return regularCost;
             }
 
-            // Mode-decision costs already contain the selected coefficient syntax but not the block's skip
-            // or IBC choice. When regular intra skips, replace its empty-coefficient rate with skip syntax.
             int skipContext = Av1TileWriter.GetSkipContext(macroBlock);
-            int regularRateAdjustment = writer.GetUseIntraBlockCopyCost(false) +
-                writer.GetSkipCost(modeInfo.Block.Skip, skipContext);
-
-            if (modeInfo.Block.Skip)
-            {
-                regularRateAdjustment -= regularEmptyTransformRate;
-            }
-
-            long bestCost = regularModeCost + Av1RateDistortion.GetCost(this.rateMultiplier, regularRateAdjustment, 0);
+            long bestCost = regularCost;
             bool hasSelectedCandidate = false;
             bool selectedSkip = false;
             Av1MotionVector selectedVector = default;
@@ -369,7 +358,7 @@ internal static partial class Av1IntraSuperblockEncoder
 
             if (!hasSelectedCandidate)
             {
-                return;
+                return bestCost;
             }
 
             // Only the winning vector is now visible to later coding blocks. This single publication keeps
@@ -440,6 +429,7 @@ internal static partial class Av1IntraSuperblockEncoder
             block.PredictionUnit.ChromaFromLumaSigns = 0;
             paletteInfo = default;
             this.picture.SetDisplacementVector(modeInfoPosition, selectedVector);
+            return bestCost;
         }
 
         private void EvaluateIntraBlockCopyPlane(

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

@@ -32,19 +32,21 @@ internal sealed class Av1EncoderPictureBuffer : IDisposable
     /// <param name="frameHeader">The frame header defining dimensions and tiles.</param>
     /// <param name="width">The visible luma width.</param>
     /// <param name="height">The visible luma height.</param>
+    /// <param name="disallow4x4AllFrames">Whether each allocated mode-information value represents an 8x8 region.</param>
     public Av1EncoderPictureBuffer(
         Configuration configuration,
         ObuSequenceHeader sequenceHeader,
         ObuFrameHeader frameHeader,
         int width,
-        int height)
+        int height,
+        bool disallow4x4AllFrames)
     {
         const int ContextAlignmentLog2 = Av1Constants.MaxSuperBlockSizeLog2 - Av1Constants.ModeInfoSizeLog2;
         this.modeInfo = new Av1EncoderModeInfoBuffer(
             configuration,
             width,
             height,
-            disallow4x4AllFrames: true);
+            disallow4x4AllFrames);
 
         int alignedModeInfoRowCount = Av1Math.AlignPowerOf2(this.modeInfo.ModeInfoRowCount, ContextAlignmentLog2);
         int lumaLeftLength = alignedModeInfoRowCount;

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1TileWriter.BlockEncoding.cs

@@ -15,6 +15,24 @@ internal partial class Av1TileWriter
     /// </summary>
     internal interface IBlockEncodingHandler
     {
+        /// <summary>
+        /// Selects the partition used for the current tree node.
+        /// </summary>
+        /// <param name="writer">The live tile symbol encoder.</param>
+        /// <param name="macroBlock">The tile-local macroblock state.</param>
+        /// <param name="blockOrigin">The absolute luma-sample origin.</param>
+        /// <param name="tileIndex">The zero-based tile index.</param>
+        /// <param name="blockSize">The current square partition size.</param>
+        /// <param name="preparedPartition">The partition retained before live analysis.</param>
+        /// <returns>The partition to encode.</returns>
+        Av1PartitionType SelectPartition(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Point blockOrigin,
+            ushort tileIndex,
+            Av1BlockSize blockSize,
+            Av1PartitionType preparedPartition);
+
         /// <summary>
         /// Encodes one final block against the current reconstructed neighbors and live tile probabilities.
         /// </summary>
@@ -37,6 +55,17 @@ internal partial class Av1TileWriter
 
     private readonly struct PrecomputedBlockEncodingHandler : IBlockEncodingHandler
     {
+        /// <inheritdoc/>
+        public Av1PartitionType SelectPartition(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Point blockOrigin,
+            ushort tileIndex,
+            Av1BlockSize blockSize,
+            Av1PartitionType preparedPartition)
+            => preparedPartition;
+
+        /// <inheritdoc/>
         public void EncodeBlock(
             Av1SymbolEncoder writer,
             Av1MacroBlockD macroBlock,

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

@@ -149,7 +149,19 @@ internal partial class Av1TileWriter
             return;
         }
 
-        Av1PartitionType partition = (Av1PartitionType)superblock.CodingUnitPartitionTypes[partitionIndex++];
+        int currentPartitionIndex = partitionIndex++;
+        Av1PartitionType preparedPartition =
+            (Av1PartitionType)superblock.CodingUnitPartitionTypes[currentPartitionIndex];
+
+        Av1PartitionType partition = blockEncoder.SelectPartition(
+            writer,
+            entropyCodingContext.MacroBlock,
+            blockOrigin,
+            tileIndex,
+            blockSize,
+            preparedPartition);
+
+        superblock.CodingUnitPartitionTypes[currentPartitionIndex] = (byte)partition;
         Av1BlockSize subSize = partition.GetBlockSubSize(blockSize);
         int halfBlockSize = blockSize.GetWidth() >> 1;
         int quarterBlockSize = blockSize.GetWidth() >> 2;
@@ -232,54 +244,89 @@ internal partial class Av1TileWriter
 
                 break;
             case Av1PartitionType.Split:
-                WritePartitionTree(
-                    pcs,
-                    entropyCodingContext,
-                    writer,
-                    superblock,
-                    coefficientBuffer,
-                    tileIndex,
-                    subSize,
-                    blockOrigin,
-                    ref partitionIndex,
-                    ref finalBlockIndex,
-                    ref blockEncoder);
-                WritePartitionTree(
-                    pcs,
-                    entropyCodingContext,
-                    writer,
-                    superblock,
-                    coefficientBuffer,
-                    tileIndex,
-                    subSize,
-                    blockOrigin + new Size(halfBlockSize, 0),
-                    ref partitionIndex,
-                    ref finalBlockIndex,
-                    ref blockEncoder);
-                WritePartitionTree(
-                    pcs,
-                    entropyCodingContext,
-                    writer,
-                    superblock,
-                    coefficientBuffer,
-                    tileIndex,
-                    subSize,
-                    blockOrigin + new Size(0, halfBlockSize),
-                    ref partitionIndex,
-                    ref finalBlockIndex,
-                    ref blockEncoder);
-                WritePartitionTree(
-                    pcs,
-                    entropyCodingContext,
-                    writer,
-                    superblock,
-                    coefficientBuffer,
-                    tileIndex,
-                    subSize,
-                    blockOrigin + new Size(halfBlockSize, halfBlockSize),
-                    ref partitionIndex,
-                    ref finalBlockIndex,
-                    ref blockEncoder);
+                if (blockSize == Av1BlockSize.Block8x8)
+                {
+                    // A split 8x8 node terminates in four 4x4 coding blocks. AV1 does not carry another
+                    // partition symbol at that size, so the children are final blocks rather than tree nodes.
+                    for (int childIndex = 0; childIndex < 4; childIndex++)
+                    {
+                        Point childOrigin = blockOrigin + new Size(
+                            (childIndex & 1) * halfBlockSize,
+                            (childIndex >> 1) * halfBlockSize);
+
+                        Point childModeInfoPosition = childOrigin >> Av1Constants.ModeInfoSizeLog2;
+                        if (childModeInfoPosition.Y >= common.ModeInfoRowCount ||
+                            childModeInfoPosition.X >= common.ModeInfoColumnCount)
+                        {
+                            continue;
+                        }
+
+                        WriteFinalBlock(
+                            pcs,
+                            entropyCodingContext,
+                            writer,
+                            superblock,
+                            coefficientBuffer,
+                            tileIndex,
+                            childOrigin,
+                            ref finalBlockIndex,
+                            ref blockEncoder);
+                    }
+                }
+                else
+                {
+                    WritePartitionTree(
+                        pcs,
+                        entropyCodingContext,
+                        writer,
+                        superblock,
+                        coefficientBuffer,
+                        tileIndex,
+                        subSize,
+                        blockOrigin,
+                        ref partitionIndex,
+                        ref finalBlockIndex,
+                        ref blockEncoder);
+
+                    WritePartitionTree(
+                        pcs,
+                        entropyCodingContext,
+                        writer,
+                        superblock,
+                        coefficientBuffer,
+                        tileIndex,
+                        subSize,
+                        blockOrigin + new Size(halfBlockSize, 0),
+                        ref partitionIndex,
+                        ref finalBlockIndex,
+                        ref blockEncoder);
+
+                    WritePartitionTree(
+                        pcs,
+                        entropyCodingContext,
+                        writer,
+                        superblock,
+                        coefficientBuffer,
+                        tileIndex,
+                        subSize,
+                        blockOrigin + new Size(0, halfBlockSize),
+                        ref partitionIndex,
+                        ref finalBlockIndex,
+                        ref blockEncoder);
+
+                    WritePartitionTree(
+                        pcs,
+                        entropyCodingContext,
+                        writer,
+                        superblock,
+                        coefficientBuffer,
+                        tileIndex,
+                        subSize,
+                        blockOrigin + new Size(halfBlockSize, halfBlockSize),
+                        ref partitionIndex,
+                        ref finalBlockIndex,
+                        ref blockEncoder);
+                }
 
                 break;
             case Av1PartitionType.HorizontalA:
@@ -592,6 +639,47 @@ internal partial class Av1TileWriter
             edgeMask);
     }
 
+    /// <summary>
+    /// Gets the partition-symbol rate from the above and left contexts available at a block origin.
+    /// </summary>
+    /// <param name="pcs">The picture coding state.</param>
+    /// <param name="writer">The live tile symbol encoder.</param>
+    /// <param name="blockSize">The square parent block size.</param>
+    /// <param name="partitionType">The partition type to measure.</param>
+    /// <param name="blockOrigin">The block origin in samples.</param>
+    /// <param name="partitionContexts">The partition neighbor arrays for the tile.</param>
+    /// <returns>The rate cost in 1/512-bit units.</returns>
+    public static int GetPartitionCost(
+        Av1PictureControlSet pcs,
+        Av1SymbolEncoder writer,
+        Av1BlockSize blockSize,
+        Av1PartitionType partitionType,
+        Point blockOrigin,
+        Av1NeighborArrayUnit<Av1PartitionContext> partitionContexts)
+    {
+        int context = GetPartitionContext(
+            pcs,
+            blockSize,
+            blockOrigin,
+            partitionContexts,
+            out bool hasRows,
+            out bool hasColumns);
+
+        if (!hasRows && !hasColumns)
+        {
+            return 0;
+        }
+
+        if (hasRows && hasColumns)
+        {
+            return writer.GetPartitionTypeCost(partitionType, context);
+        }
+
+        return !hasRows
+            ? writer.GetSplitOrHorizontalCost(partitionType, blockSize, context)
+            : writer.GetSplitOrVerticalCost(partitionType, blockSize, context);
+    }
+
     /// <summary>
     /// Writes a partition symbol using the above and left partition contexts available at a block origin.
     /// </summary>
@@ -616,33 +704,13 @@ internal partial class Av1TileWriter
             return;
         }
 
-        int halfBlockModeInfoCount = blockSize.Get4x4WideCount() >> 1;
-        Point modeInfoPosition = blockOrigin >> Av1Constants.ModeInfoSizeLog2;
-        bool has_rows = modeInfoPosition.Y + halfBlockModeInfoCount < pcs.Parent.Common.ModeInfoRowCount;
-        bool has_cols = modeInfoPosition.X + halfBlockModeInfoCount < pcs.Parent.Common.ModeInfoColumnCount;
-
-        int partition_context_left_neighbor_index = partition_context_na.GetLeftIndex(blockOrigin);
-        int partition_context_top_neighbor_index = partition_context_na.GetTopIndex(blockOrigin);
-
-        int context_index = 0;
-
-        byte above_ctx =
-            (byte)(partition_context_na.Top[partition_context_top_neighbor_index].Above == byte.MaxValue
-            ? 0
-            : partition_context_na.Top[partition_context_top_neighbor_index].Above);
-        byte left_ctx =
-            (byte)(partition_context_na.Left[partition_context_left_neighbor_index].Left == byte.MaxValue
-            ? 0
-            : partition_context_na.Left[partition_context_left_neighbor_index].Left);
-
-        int blockSizeLog2 = blockSize.Get4x4WidthLog2() - 1;
-        int above = (above_ctx >> blockSizeLog2) & 1, left = (left_ctx >> blockSizeLog2) & 1;
-
-        Guard.IsTrue(blockSize.Get4x4WidthLog2() == blockSize.Get4x4HeightLog2(), nameof(blockSize), "Blocks need to be square.");
-        Guard.IsTrue(blockSizeLog2 >= 0, nameof(blockSizeLog2), "bsl needs to be a positive integer.");
-
-        // Each square block-size level owns four contexts selected by the current split bit of its neighbors.
-        context_index = ((left * 2) + above) + (blockSizeLog2 * Av1Constants.PartitionProbabilitySet);
+        int context_index = GetPartitionContext(
+            pcs,
+            blockSize,
+            blockOrigin,
+            partition_context_na,
+            out bool has_rows,
+            out bool has_cols);
 
         if (!has_rows && !has_cols)
         {
@@ -666,6 +734,39 @@ internal partial class Av1TileWriter
         return;
     }
 
+    private static int GetPartitionContext(
+        Av1PictureControlSet pcs,
+        Av1BlockSize blockSize,
+        Point blockOrigin,
+        Av1NeighborArrayUnit<Av1PartitionContext> partitionContexts,
+        out bool hasRows,
+        out bool hasColumns)
+    {
+        int halfBlockModeInfoCount = blockSize.Get4x4WideCount() >> 1;
+        Point modeInfoPosition = blockOrigin >> Av1Constants.ModeInfoSizeLog2;
+        hasRows = modeInfoPosition.Y + halfBlockModeInfoCount < pcs.Parent.Common.ModeInfoRowCount;
+        hasColumns = modeInfoPosition.X + halfBlockModeInfoCount < pcs.Parent.Common.ModeInfoColumnCount;
+        int leftIndex = partitionContexts.GetLeftIndex(blockOrigin);
+        int topIndex = partitionContexts.GetTopIndex(blockOrigin);
+        byte aboveContext = partitionContexts.Top[topIndex].Above == byte.MaxValue
+            ? (byte)0
+            : partitionContexts.Top[topIndex].Above;
+
+        byte leftContext = partitionContexts.Left[leftIndex].Left == byte.MaxValue
+            ? (byte)0
+            : partitionContexts.Left[leftIndex].Left;
+
+        int blockSizeLog2 = blockSize.Get4x4WidthLog2() - 1;
+        int above = (aboveContext >> blockSizeLog2) & 1;
+        int left = (leftContext >> blockSizeLog2) & 1;
+
+        Guard.IsTrue(blockSize.Get4x4WidthLog2() == blockSize.Get4x4HeightLog2(), nameof(blockSize), "Blocks need to be square.");
+        Guard.IsTrue(blockSizeLog2 >= 0, nameof(blockSizeLog2), "bsl needs to be a positive integer.");
+
+        // Each square block-size level owns four contexts selected by the current split bit of its neighbors.
+        return ((left * 2) + above) + (blockSizeLog2 * Av1Constants.PartitionProbabilitySet);
+    }
+
     /// <summary>
     /// Writes the segmentation, prediction, transform, coefficient, and filter syntax for one final coding block.
     /// </summary>
@@ -1900,7 +2001,7 @@ internal partial class Av1TileWriter
     /// <param name="subsamplingX">The horizontal chroma subsampling shift.</param>
     /// <param name="subsamplingY">The vertical chroma subsampling shift.</param>
     /// <returns>The aligned origin in chroma samples.</returns>
-    private static Point GetChromaBlockOrigin(Point lumaOrigin, int subsamplingX, int subsamplingY)
+    public static Point GetChromaBlockOrigin(Point lumaOrigin, int subsamplingX, int subsamplingY)
         => new(
             (lumaOrigin.X >> (Av1Constants.ModeInfoSizeLog2 + subsamplingX)) << Av1Constants.ModeInfoSizeLog2,
             (lumaOrigin.Y >> (Av1Constants.ModeInfoSizeLog2 + subsamplingY)) << Av1Constants.ModeInfoSizeLog2);

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

@@ -368,7 +368,8 @@ public class Av1CoefficientsEntropyTests
             sequenceHeader,
             frameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         Av1PictureControlSet picture = pictureBuffer.Picture;
         Av1NeighborArrayUnit<Av1EncoderPaletteInfo> paletteContexts = Assert.Single(picture.PaletteContexts);

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

@@ -153,6 +153,69 @@ public class Av1EncoderFrameTests
         }
     }
 
+    [Fact]
+    public void EncodeEffortNineSelectsSubEightPartition()
+    {
+        const int Size = 16;
+        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++)
+            {
+                // The bottom-right 8x8 uses horizontal prediction on its left half and vertical prediction
+                // on its right half. Twelve source values keep a parent palette from reproducing both halves.
+                byte value;
+                if (x < 8 && y < 8)
+                {
+                    value = 128;
+                }
+                else if (y < 8)
+                {
+                    value = (byte)(16 + ((x - 8) * 20));
+                }
+                else
+                {
+                    value = x < 12
+                        ? (byte)(176 + ((y - 8) * 9))
+                        : (byte)(16 + ((x - 8) * 20));
+                }
+
+                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);
+        Point[] leafPositions =
+        [
+            new(2, 2),
+            new(3, 2),
+            new(2, 3),
+            new(3, 3)
+        ];
+
+        foreach (Point leafPosition in leafPositions)
+        {
+            Assert.Equal(
+                Av1BlockSize.Block4x8,
+                frameInfo.GetModeInfoAt(leafPosition).BlockSize);
+        }
+
+        Assert.Equal(new Size(Size, Size), decoded.Size);
+    }
+
     [Theory]
     [InlineData(EightBit)]
     [InlineData(TenBit)]

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

@@ -99,7 +99,8 @@ public class Av1EncoderModeInfoBufferTests
             sequenceHeader,
             frameHeader,
             Width,
-            Height))
+            Height,
+            disallow4x4AllFrames: true))
         {
             allocations = allocator.AllocationLog.ToArray();
             Assert.Equal(2, allocations.Length);

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

@@ -100,7 +100,8 @@ public class Av1IntraBlockCopyTests
             sequenceHeader,
             frameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         Av1PictureControlSet picture = buffer.Picture;
         Point candidatePosition = new(80, 12);
@@ -161,7 +162,8 @@ public class Av1IntraBlockCopyTests
             sequenceHeader,
             frameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         Av1PictureControlSet picture = buffer.Picture;
         Point modeInfoPosition = new(80, 0);
@@ -254,7 +256,8 @@ public class Av1IntraBlockCopyTests
             sequenceHeader,
             frameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderFrameBuffer<byte> source = new(
             Configuration.Default,
@@ -339,7 +342,8 @@ public class Av1IntraBlockCopyTests
             sequenceHeader,
             frameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderFrameBuffer<byte> source = new(
             Configuration.Default,
@@ -434,7 +438,8 @@ public class Av1IntraBlockCopyTests
             sequenceHeader,
             frameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderFrameBuffer<byte> source = new(
             Configuration.Default,

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

@@ -244,7 +244,8 @@ public class Av1IntraSuperblockEncoderTests
             picture.Sequence.SequenceHeader,
             picture.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer tileCoefficients = new(
             Configuration.Default,
@@ -320,7 +321,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         Av1PictureControlSet picture = pictureBuffer.Picture;
         using Av1EncoderCoefficientBuffer coefficients = new(
@@ -393,7 +395,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer liveCoefficients = new(
             Configuration.Default,
@@ -517,7 +520,8 @@ public class Av1IntraSuperblockEncoderTests
             picture.Sequence.SequenceHeader,
             picture.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer tileCoefficients = new(
             Configuration.Default,
@@ -572,7 +576,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -664,7 +669,8 @@ public class Av1IntraSuperblockEncoderTests
                 sequenceHeader,
                 frameHeader,
                 Width,
-                Height);
+                Height,
+                disallow4x4AllFrames: true);
 
             using Av1EncoderCoefficientBuffer coefficients = new(
                 Configuration.Default,
@@ -921,7 +927,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -1133,7 +1140,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -1246,7 +1254,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -1406,7 +1415,8 @@ public class Av1IntraSuperblockEncoderTests
             pilotTemplate.Sequence.SequenceHeader,
             pilotTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer pilotCoefficients = new(
             Configuration.Default,
@@ -1507,7 +1517,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -1747,7 +1758,8 @@ public class Av1IntraSuperblockEncoderTests
             pilotTemplate.Sequence.SequenceHeader,
             pilotTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer pilotCoefficients = new(
             Configuration.Default,
@@ -1903,7 +1915,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -2101,7 +2114,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -2243,7 +2257,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -2370,7 +2385,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -2480,7 +2496,8 @@ public class Av1IntraSuperblockEncoderTests
             sequenceHeader,
             frameHeader,
             Width,
-            Height);
+            Height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -2738,7 +2755,8 @@ public class Av1IntraSuperblockEncoderTests
             pictureTemplate.Sequence.SequenceHeader,
             pictureTemplate.Parent.FrameHeader,
             width,
-            height);
+            height,
+            disallow4x4AllFrames: true);
 
         using Av1EncoderCoefficientBuffer coefficients = new(
             Configuration.Default,
@@ -3035,6 +3053,16 @@ public class Av1IntraSuperblockEncoderTests
         /// </summary>
         public int Count { get; private set; }
 
+        /// <inheritdoc/>
+        public readonly Av1PartitionType SelectPartition(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Point blockOrigin,
+            ushort tileIndex,
+            Av1BlockSize blockSize,
+            Av1PartitionType preparedPartition)
+            => preparedPartition;
+
         /// <inheritdoc/>
         public void EncodeBlock(
             Av1SymbolEncoder writer,
@@ -3100,6 +3128,16 @@ public class Av1IntraSuperblockEncoderTests
         /// </summary>
         public int Count { get; private set; }
 
+        /// <inheritdoc/>
+        public readonly Av1PartitionType SelectPartition(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Point blockOrigin,
+            ushort tileIndex,
+            Av1BlockSize blockSize,
+            Av1PartitionType preparedPartition)
+            => preparedPartition;
+
         /// <inheritdoc/>
         public void EncodeBlock(
             Av1SymbolEncoder writer,