Complete AV1 128x128 partition search

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

@@ -18,29 +18,35 @@ internal readonly ref struct Av1EncoderModeDecisionWorkspace<TSample>
     /// <summary>
     /// The largest coding-block dimension evaluated directly by the current partition search.
     /// </summary>
-    public const int MaximumBlockDimension = 64;
+    public const int MaximumBlockDimension = 128;
 
     /// <summary>
     /// The maximum number of samples in one directly evaluated coding block.
     /// </summary>
     public const int MaximumSampleCount = MaximumBlockDimension * MaximumBlockDimension;
 
+    /// <summary>
+    /// The maximum number of samples in one AV1 transform.
+    /// </summary>
+    public const int MaximumTransformSampleCount =
+        Av1Constants.MaxTransformSize * Av1Constants.MaxTransformSize;
+
     /// <summary>
     /// The number of 4x4 transform blocks covering one 8x8 coding block.
     /// </summary>
     public const int CandidateTransformBlockCount = 4;
 
     /// <summary>
-    /// The maximum number of transform states needed while evaluating both chroma planes of one 64x64 block.
+    /// The maximum number of transform states needed while evaluating both chroma planes of one 128x128 block.
     /// </summary>
-    public const int MaximumCandidateTransformBlockCount = 8;
+    public const int MaximumCandidateTransformBlockCount = 32;
 
     /// <summary>
     /// The required workspace length in signed-integer storage elements.
     /// </summary>
     public const int StorageLength = TransientStorageOffset + Av1EncoderPaletteWorkspace<ushort>.StorageLength;
 
-    private const int ReferenceBufferLength = (2 * MaximumBlockDimension) + 1;
+    private const int ReferenceBufferLength = (2 * Av1Constants.MaxTransformSize) + 1;
     private const int ReferenceBufferCount = 4;
     private const int ReferenceStorageLength = ReferenceBufferCount * ReferenceBufferLength * sizeof(ushort) / sizeof(int);
     private const int CandidateSampleStorageOffset = ReferenceStorageLength;
@@ -80,7 +86,7 @@ internal readonly ref struct Av1EncoderModeDecisionWorkspace<TSample>
     /// Gets the temporary prediction span shared by mutually exclusive mode searches.
     /// </summary>
     public Span<TSample> Prediction
-        => MemoryMarshal.Cast<int, TSample>(this.storage[TransientStorageOffset..])[..MaximumSampleCount];
+        => MemoryMarshal.Cast<int, TSample>(this.storage[TransientStorageOffset..])[..MaximumTransformSampleCount];
 
     /// <summary>
     /// Gets the temporary residual span shared by mutually exclusive mode searches.
@@ -88,8 +94,8 @@ internal readonly ref struct Av1EncoderModeDecisionWorkspace<TSample>
     public Span<short> Residual
         => MemoryMarshal.Cast<int, short>(
             this.storage.Slice(
-                TransientStorageOffset + (MaximumSampleCount * sizeof(ushort) / sizeof(int)),
+                TransientStorageOffset + (MaximumTransformSampleCount * sizeof(ushort) / sizeof(int)),
-                MaximumSampleCount * sizeof(short) / sizeof(int)));
+                MaximumTransformSampleCount * sizeof(short) / sizeof(int)));
 
     /// <summary>
     /// Gets the fixed-stride subsampled luma values used by chroma-from-luma mode search.
@@ -182,7 +188,8 @@ internal readonly ref struct Av1EncoderPaletteWorkspace<TSample>
     /// </summary>
     public const int StorageLength = ColorCacheOffset + ColorCacheStorageLength;
 
-    private const int MaximumSampleCount = Av1EncoderModeDecisionWorkspace<ushort>.MaximumSampleCount;
+    private const int MaximumBlockDimension = 64;
+    private const int MaximumSampleCount = MaximumBlockDimension * MaximumBlockDimension;
     private const int PlaneShortStorageLength = MaximumSampleCount * sizeof(short) / sizeof(int);
     private const int PlaneSampleStorageLength = MaximumSampleCount * sizeof(ushort) / sizeof(int);
     private const int PlaneByteStorageLength = MaximumSampleCount / sizeof(int);

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

@@ -89,7 +89,7 @@ internal static class Av1FrameEncoder
             FrameHeightBits = height > 1 ? Av1Math.MostSignificantBit((uint)(height - 1)) + 1 : 1,
             MaxFrameWidth = width,
             MaxFrameHeight = height,
-            Use128x128Superblock = false,
+            Use128x128Superblock = effort == 10 && width >= 128 && height >= 128,
             ForceScreenContentTools = 2,
             ForceIntegerMotionVector = 2,
             EnableFilterIntra = effort >= 4,

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

@@ -217,8 +217,9 @@ internal static partial class Av1IntraSuperblockEncoder
             };
 
             bool hasLumaPalette = paletteInfo.PaletteSizes[0] != 0;
-            int paletteDisabledCost = blockSize >= Av1BlockSize.Block8x8 &&
+            int paletteDisabledCost = Av1TileWriter.IsPaletteAllowed(
-                this.picture.Parent.FrameHeader.AllowScreenContentTools
+                this.picture.Parent.FrameHeader.AllowScreenContentTools,
+                blockSize)
                 ? writer.GetPaletteUvModeCost(false, hasLumaPalette)
                 : 0;
 
@@ -575,6 +576,9 @@ internal static partial class Av1IntraSuperblockEncoder
             int transformColumnCount = blockWidth / transformWidth;
             int transformRowCount = blockHeight / transformHeight;
             int transformBlockCount = transformColumnCount * transformRowCount;
+            Av1BlockSize maximumUnitBlockSize =
+                Av1BlockSize.Block64x64.GetSubsampled(colorConfig.SubSamplingX, colorConfig.SubSamplingY);
+
             Span<TSample> candidateBlueReconstruction =
                 workspace.GetCandidateReconstruction(0)[..blockSampleCount];
 
@@ -619,7 +623,9 @@ internal static partial class Av1IntraSuperblockEncoder
             Buffer2DRegion<TSample> blueReconstruction = this.reconstruction.GetPlane(Av1Plane.U);
             Buffer2DRegion<TSample> redReconstruction = this.reconstruction.GetPlane(Av1Plane.V);
             bool hasLumaPalette = paletteInfo.PaletteSizes[0] != 0;
-            int paletteDisabledCost = this.picture.Parent.FrameHeader.AllowScreenContentTools
+            int paletteDisabledCost = Av1TileWriter.IsPaletteAllowed(
+                this.picture.Parent.FrameHeader.AllowScreenContentTools,
+                blockSize)
                 ? writer.GetPaletteUvModeCost(false, hasLumaPalette)
                 : 0;
 
@@ -662,7 +668,7 @@ internal static partial class Av1IntraSuperblockEncoder
                 redNeighbors.Top.Slice(redTopIndex, contextWidth).CopyTo(redTopContexts);
                 redNeighbors.Left.Slice(redLeftIndex, contextHeight).CopyTo(redLeftContexts);
                 Av1PredictionMode predictionMode = chromaMode.ToLumaMode();
-                long distortion = this.GetTiledChromaPlaneCost(
+                long distortion = this.GetTiledPlaneCost(
                     writer,
                     macroBlock,
                     lumaOrigin,
@@ -670,6 +676,7 @@ internal static partial class Av1IntraSuperblockEncoder
                     blockSize,
                     chromaBlockSize,
                     transformSize,
+                    maximumUnitBlockSize,
                     subsamplingX,
                     subsamplingY,
                     lumaMode,
@@ -685,7 +692,7 @@ internal static partial class Av1IntraSuperblockEncoder
                     blueLeftContexts,
                     out int blueRate);
 
-                distortion += this.GetTiledChromaPlaneCost(
+                distortion += this.GetTiledPlaneCost(
                     writer,
                     macroBlock,
                     lumaOrigin,
@@ -693,6 +700,7 @@ internal static partial class Av1IntraSuperblockEncoder
                     blockSize,
                     chromaBlockSize,
                     transformSize,
+                    maximumUnitBlockSize,
                     subsamplingX,
                     subsamplingY,
                     lumaMode,
@@ -761,7 +769,7 @@ internal static partial class Av1IntraSuperblockEncoder
             return bestMode;
         }
 
-        private long GetTiledChromaPlaneCost(
+        private long GetTiledPlaneCost(
             Av1SymbolEncoder writer,
             Av1MacroBlockD macroBlock,
             Point lumaOrigin,
@@ -769,6 +777,7 @@ internal static partial class Av1IntraSuperblockEncoder
             Av1BlockSize blockSize,
             Av1BlockSize chromaBlockSize,
             Av1TransformSize transformSize,
+            Av1BlockSize maximumUnitBlockSize,
             int subsamplingX,
             int subsamplingY,
             Av1PredictionMode lumaMode,
@@ -794,12 +803,17 @@ internal static partial class Av1IntraSuperblockEncoder
             int transformSampleCount = transformSize.GetSize2d();
             int transformWidth4x4 = transformSize.Get4x4WideCount();
             int transformHeight4x4 = transformSize.Get4x4HighCount();
+            int maximumUnitWidth = Math.Min(maximumUnitBlockSize.GetWidth(), blockWidth);
+            int maximumUnitHeight = Math.Min(maximumUnitBlockSize.GetHeight(), blockHeight);
             Av1TransformType transformType = Av1SymbolContextHelper.GetDefaultIntraTransformType(
                 predictionMode,
                 transformSize,
                 this.picture.Parent.FrameHeader.UseReducedTransformSet);
 
-            Av1ComponentType componentType = Av1ComponentType.Chroma;
+            Av1ComponentType componentType = plane == Av1Plane.Y
+                ? Av1ComponentType.Luminance
+                : Av1ComponentType.Chroma;
+
             Span<TSample> prediction = workspace.Prediction[..transformSampleCount];
             Span<short> residual = workspace.Residual[..transformSampleCount];
             Span<TSample> aboveStorage = workspace.GetReferenceSamples(0);
@@ -808,104 +822,115 @@ internal static partial class Av1IntraSuperblockEncoder
             int transformIndex = 0;
             long distortion = 0;
             rate = 0;
-            for (int transformRow = 0; transformRow < blockHeight / transformHeight; transformRow++)
+
+            // Residual syntax completes each bounded 64x64 luma region, scaled for chroma, before
+            // moving to the next region. Candidate coefficients and states must retain that exact order.
+            for (int regionRow = 0; regionRow < blockHeight; regionRow += maximumUnitHeight)
             {
-                int rowOffset = transformRow * transformHeight;
+                int unitBottom = Math.Min(regionRow + maximumUnitHeight, blockHeight);
-                for (int transformColumn = 0; transformColumn < blockWidth / transformWidth; transformColumn++)
+                for (int regionColumn = 0; regionColumn < blockWidth; regionColumn += maximumUnitWidth)
                 {
-                    int columnOffset = transformColumn * transformWidth;
+                    int unitRight = Math.Min(regionColumn + maximumUnitWidth, blockWidth);
-                    int reconstructionOffset = (rowOffset * blockWidth) + columnOffset;
+                    for (int rowOffset = regionRow; rowOffset < unitBottom; rowOffset += transformHeight)
-                    Point transformOrigin = chromaOrigin + new Size(columnOffset, rowOffset);
+                    {
-                    this.PrepareTransformReferenceSamples(
+                        int transformRow = rowOffset / transformHeight;
-                        reconstruction,
+                        for (int columnOffset = regionColumn; columnOffset < unitRight; columnOffset += transformWidth)
-                        lumaOrigin,
+                        {
-                        chromaOrigin,
+                            int transformColumn = columnOffset / transformWidth;
-                        blockSize,
+                            int reconstructionOffset = (rowOffset * blockWidth) + columnOffset;
-                        macroBlock,
+                            Point transformOrigin = chromaOrigin + new Size(columnOffset, rowOffset);
-                        transformRow,
+                            this.PrepareTransformReferenceSamples(
-                        transformColumn,
+                                reconstruction,
-                        blockWidth,
+                                lumaOrigin,
-                        transformSize,
+                                chromaOrigin,
-                        subsamplingX,
+                                blockSize,
-                        subsamplingY,
+                                macroBlock,
-                        candidateReconstruction,
+                                transformRow,
-                        aboveStorage,
+                                transformColumn,
-                        leftStorage,
+                                blockWidth,
-                        out bool hasLeft,
+                                transformSize,
-                        out bool hasAbove);
+                                subsamplingX,
-
+                                subsamplingY,
-                    TOperator.PrepareIntra(
+                                candidateReconstruction,
-                        this.blockWorkspace,
+                                aboveStorage,
-                        source,
+                                leftStorage,
-                        transformOrigin,
+                                out bool hasLeft,
-                        prediction,
+                                out bool hasAbove);
-                        aboveStorage.Slice(1, transformWidth * 2),
+
-                        leftStorage.Slice(1, transformHeight * 2),
+                            TOperator.PrepareIntra(
-                        hasLeft,
+                                this.blockWorkspace,
-                        hasAbove,
+                                source,
-                        predictionMode,
+                                transformOrigin,
-                        angleDelta,
+                                prediction,
-                        residual,
+                                aboveStorage.Slice(1, transformWidth * 2),
-                        transformSize,
+                                leftStorage.Slice(1, transformHeight * 2),
-                        this.bitDepth);
+                                hasLeft,
-
+                                hasAbove,
-                    Av1TransformBlockContext blockContext = Av1TileWriter.GetTransformBlockContexts(
+                                predictionMode,
-                        componentType,
+                                angleDelta,
-                        topContexts.Slice(transformColumn * transformWidth4x4, transformWidth4x4),
+                                residual,
-                        leftContexts.Slice(transformRow * transformHeight4x4, transformHeight4x4),
+                                transformSize,
-                        chromaBlockSize,
+                                this.bitDepth);
-                        transformSize);
+
-
+                            Av1TransformBlockContext blockContext = Av1TileWriter.GetTransformBlockContexts(
-                    Span<int> transformCoefficients = candidateCoefficients.Slice(
+                                componentType,
-                        coefficientOffset,
+                                topContexts.Slice(transformColumn * transformWidth4x4, transformWidth4x4),
-                        transformSampleCount);
+                                leftContexts.Slice(transformRow * transformHeight4x4, transformHeight4x4),
-
+                                chromaBlockSize,
-                    ref Av1EncoderTransformBlockState state = ref candidateStates[transformIndex++];
+                                transformSize);
-                    distortion += TOperator.EncodePredictionCandidate(
+
-                        this.blockWorkspace,
+                            Span<int> transformCoefficients = candidateCoefficients.Slice(
-                        source,
+                                coefficientOffset,
-                        transformOrigin,
+                                transformSampleCount);
-                        prediction,
+
-                        residual,
+                            ref Av1EncoderTransformBlockState state = ref candidateStates[transformIndex++];
-                        candidateReconstruction[reconstructionOffset..],
+                            distortion += TOperator.EncodePredictionCandidate(
-                        blockWidth,
+                                this.blockWorkspace,
-                        transformCoefficients,
+                                source,
-                        transformSize,
+                                transformOrigin,
-                        transformType,
+                                prediction,
-                        plane,
+                                residual,
-                        this.quantization.QIndex[0],
+                                candidateReconstruction[reconstructionOffset..],
-                        this.quantization.DeltaQDc[(int)plane],
+                                blockWidth,
-                        this.quantization.DeltaQAc[(int)plane],
+                                transformCoefficients,
-                        this.bitDepth,
+                                transformSize,
-                        ref state);
+                                transformType,
-
+                                plane,
-                    rate += writer.GetCoefficientCost(
+                                this.quantization.QIndex[0],
-                        transformSize,
+                                this.quantization.DeltaQDc[(int)plane],
-                        transformType,
+                                this.quantization.DeltaQAc[(int)plane],
-                        lumaMode,
+                                this.bitDepth,
-                        transformCoefficients,
+                                ref state);
-                        componentType,
+
-                        blockContext,
+                            rate += writer.GetCoefficientCost(
-                        state.EndOfBlock,
+                                transformSize,
-                        this.picture.Parent.FrameHeader.UseReducedTransformSet,
+                                transformType,
-                        Av1FilterIntraMode.AllFilterIntraModes,
+                                lumaMode,
-                        usesInterTransformSet: false);
+                                transformCoefficients,
-
+                                componentType,
-                    byte coefficientContext = Av1SymbolContextHelper.GetCoefficientContext(
+                                blockContext,
-                        transformCoefficients,
+                                state.EndOfBlock,
-                        transformSize,
+                                this.picture.Parent.FrameHeader.UseReducedTransformSet,
-                        transformType,
+                                Av1FilterIntraMode.AllFilterIntraModes,
-                        state.EndOfBlock);
+                                usesInterTransformSet: false);
-
+
-                    topContexts
+                            byte coefficientContext = Av1SymbolContextHelper.GetCoefficientContext(
-                        .Slice(transformColumn * transformWidth4x4, transformWidth4x4)
+                                transformCoefficients,
-                        .Fill(coefficientContext);
+                                transformSize,
-
+                                transformType,
-                    leftContexts
+                                state.EndOfBlock);
-                        .Slice(transformRow * transformHeight4x4, transformHeight4x4)
+
-                        .Fill(coefficientContext);
+                            topContexts
-
+                                .Slice(transformColumn * transformWidth4x4, transformWidth4x4)
-                    coefficientOffset += transformSampleCount;
+                                .Fill(coefficientContext);
+
+                            leftContexts
+                                .Slice(transformRow * transformHeight4x4, transformHeight4x4)
+                                .Fill(coefficientContext);
+
+                            coefficientOffset += transformSampleCount;
+                        }
+                    }
                 }
             }
 

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

@@ -184,7 +184,8 @@ internal static partial class Av1IntraSuperblockEncoder
             Av1PartitionType preparedPartition)
         {
             bool searchPartition = blockSize is Av1BlockSize.Block8x8 or Av1BlockSize.Block16x16 ||
-                (this.effort == 10 && blockSize is Av1BlockSize.Block32x32 or Av1BlockSize.Block64x64);
+                (this.effort == 10 &&
+                    blockSize is Av1BlockSize.Block32x32 or Av1BlockSize.Block64x64 or Av1BlockSize.Block128x128);
 
             if (!searchPartition)
             {
@@ -406,6 +407,13 @@ internal static partial class Av1IntraSuperblockEncoder
                 return false;
             }
 
+            if (blockSize == Av1BlockSize.Block128x128 &&
+                partitionType is Av1PartitionType.Horizontal4 or Av1PartitionType.Vertical4)
+            {
+                // AV1 excludes 128x32 and 32x128 leaves from the 128x128 partition alphabet.
+                return false;
+            }
+
             if (this.source.IsMonochrome)
             {
                 return true;
@@ -584,13 +592,17 @@ internal static partial class Av1IntraSuperblockEncoder
 
             // Mode decision retains one state for every uniform transform tile in the coding block. The block
             // can skip coefficient syntax only when every retained transform has an empty end-of-block marker.
+            int lumaTransformSampleCount = lumaTransformSize.GetSize2d();
             int lumaTransformBlockCount =
-                (blockSize.GetWidth() * blockSize.GetHeight()) / lumaTransformSize.GetSize2d();
+                (blockSize.GetWidth() * blockSize.GetHeight()) / lumaTransformSampleCount;
+
+            int lumaStateStride =
+                lumaTransformSampleCount / Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
 
             bool lumaTransformEmpty = true;
             for (int transformIndex = 0; transformIndex < lumaTransformBlockCount; transformIndex++)
             {
-                lumaTransformEmpty &= retainedLumaStates[transformIndex].EndOfBlock == 0;
+                lumaTransformEmpty &= retainedLumaStates[transformIndex * lumaStateStride].EndOfBlock == 0;
             }
 
             if (this.source.IsMonochrome)
@@ -886,6 +898,7 @@ internal static partial class Av1IntraSuperblockEncoder
                 blockOrigin,
                 blockSize,
                 transformSize,
+                Av1BlockSize.Block64x64,
                 lumaCoefficients[lumaArea..],
                 lumaStates[(lumaArea / Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount)..]);
 
@@ -923,6 +936,9 @@ internal static partial class Av1IntraSuperblockEncoder
                 colorConfig.SubSamplingX,
                 colorConfig.SubSamplingY);
 
+            Av1BlockSize maximumChromaUnitBlockSize =
+                Av1BlockSize.Block64x64.GetSubsampled(colorConfig.SubSamplingX, colorConfig.SubSamplingY);
+
             int chromaStateIndex =
                 chromaArea / Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
 
@@ -939,6 +955,7 @@ internal static partial class Av1IntraSuperblockEncoder
                 chromaOrigin,
                 chromaBlockSize,
                 chromaTransformSize,
+                maximumChromaUnitBlockSize,
                 blueCoefficients[chromaArea..],
                 blueStates[chromaStateIndex..]);
 
@@ -947,6 +964,7 @@ internal static partial class Av1IntraSuperblockEncoder
                 chromaOrigin,
                 chromaBlockSize,
                 chromaTransformSize,
+                maximumChromaUnitBlockSize,
                 redCoefficients[chromaArea..],
                 redStates[chromaStateIndex..]);
         }
@@ -956,6 +974,7 @@ internal static partial class Av1IntraSuperblockEncoder
             Point blockOrigin,
             Av1BlockSize blockSize,
             Av1TransformSize transformSize,
+            Av1BlockSize maximumUnitBlockSize,
             ReadOnlySpan<int> coefficients,
             ReadOnlySpan<Av1EncoderTransformBlockState> states)
         {
@@ -968,32 +987,42 @@ internal static partial class Av1IntraSuperblockEncoder
             int transformWidth = transformSize.GetWidth();
             int transformHeight = transformSize.GetHeight();
             int transformSampleCount = transformSize.GetSize2d();
+            int maximumUnitWidth = Math.Min(maximumUnitBlockSize.GetWidth(), blockWidth);
+            int maximumUnitHeight = Math.Min(maximumUnitBlockSize.GetHeight(), blockHeight);
             int transformStateOffset = 0;
             int coefficientOffset = 0;
             int transformStateStride =
                 transformSampleCount / Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
 
-            // Uniform transform blocks are retained and written in raster order. Publishing that same tiling
+            // Coefficients retain AV1's bounded-region order rather than unrestricted row-major order.
-            // preserves the distinct top and left contexts consumed by the next coding block in a dry run.
+            // Publishing the same sequence pairs every state with the transform that produced it.
-            for (int row = 0; row < blockHeight; row += transformHeight)
+            for (int regionRow = 0; regionRow < blockHeight; regionRow += maximumUnitHeight)
             {
-                for (int column = 0; column < blockWidth; column += transformWidth)
+                int unitBottom = Math.Min(regionRow + maximumUnitHeight, blockHeight);
+                for (int regionColumn = 0; regionColumn < blockWidth; regionColumn += maximumUnitWidth)
                 {
-                    Av1EncoderTransformBlockState state = states[transformStateOffset];
+                    int unitRight = Math.Min(regionColumn + maximumUnitWidth, blockWidth);
-                    byte context = Av1SymbolContextHelper.GetCoefficientContext(
+                    for (int row = regionRow; row < unitBottom; row += transformHeight)
-                        coefficients[coefficientOffset..],
+                    {
-                        transformSize,
+                        for (int column = regionColumn; column < unitRight; column += transformWidth)
-                        state.TransformType,
+                        {
-                        state.EndOfBlock);
+                            Av1EncoderTransformBlockState state = states[transformStateOffset];
+                            byte context = Av1SymbolContextHelper.GetCoefficientContext(
+                                coefficients[coefficientOffset..],
+                                transformSize,
+                                state.TransformType,
+                                state.EndOfBlock);
 
-                    neighbors.UnitModeWrite(
+                            neighbors.UnitModeWrite(
-                        context,
+                                context,
-                        blockOrigin + new Size(column, row),
+                                blockOrigin + new Size(column, row),
-                        new Size(transformWidth, transformHeight),
+                                new Size(transformWidth, transformHeight),
-                        EdgeMask);
+                                EdgeMask);
 
-                    coefficientOffset += transformSampleCount;
+                            coefficientOffset += transformSampleCount;
-                    transformStateOffset += transformStateStride;
+                            transformStateOffset += transformStateStride;
+                        }
+                    }
                 }
             }
         }
@@ -1315,6 +1344,23 @@ internal static partial class Av1IntraSuperblockEncoder
 
             Buffer2DRegion<TSample> sourcePlane = this.source.GetPlane(Av1Plane.Y);
             Buffer2DRegion<TSample> reconstructionPlane = this.reconstruction.GetPlane(Av1Plane.Y);
+            if (blockWidth > transformSize.GetWidth() || blockHeight > transformSize.GetHeight())
+            {
+                return this.SelectTiledLumaMode(
+                    writer,
+                    macroBlock,
+                    blockOrigin,
+                    blockSize,
+                    tileIndex,
+                    transformSize,
+                    retainedCoefficients,
+                    retainedStates,
+                    out selectedAngleDelta,
+                    out selectedFilterIntraMode,
+                    out selectedTransformSize,
+                    out selectedCost);
+            }
+
             bool hasLeft = macroBlock.IsLeftAvailable;
             bool hasAbove = macroBlock.IsUpAvailable;
             int modeInfoRow = blockOrigin.Y >> Av1Constants.ModeInfoSizeLog2;
@@ -1438,8 +1484,9 @@ internal static partial class Av1IntraSuperblockEncoder
                 : 0;
 
             int paletteDisabledCost = 0;
-            if (blockSize >= Av1BlockSize.Block8x8 &&
+            if (Av1TileWriter.IsPaletteAllowed(
-                this.picture.Parent.FrameHeader.AllowScreenContentTools)
+                this.picture.Parent.FrameHeader.AllowScreenContentTools,
+                blockSize))
             {
                 Av1NeighborArrayUnit<Av1EncoderPaletteInfo> paletteContexts = this.picture.PaletteContexts[tileIndex];
                 int blockSizeContext = Av1TileWriter.GetPaletteBlockSizeContext(blockSize);
@@ -1698,9 +1745,9 @@ internal static partial class Av1IntraSuperblockEncoder
             }
 
             if (this.effort >= 4 &&
-                this.picture.Sequence.SequenceHeader.EnableFilterIntra &&
+                Av1TileWriter.IsFilterIntraAllowedBlockSize(
-                blockWidth <= 32 &&
+                    this.picture.Sequence.SequenceHeader.EnableFilterIntra,
-                blockHeight <= 32)
+                    blockSize))
             {
                 // Each recursive filter prediction and its source residual are independent of transform type.
                 // Prepare them once per filter mode so all legal transforms reuse the same samples.
@@ -1889,6 +1936,174 @@ internal static partial class Av1IntraSuperblockEncoder
             return bestMode;
         }
 
+        private Av1PredictionMode SelectTiledLumaMode(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Point blockOrigin,
+            Av1BlockSize blockSize,
+            ushort tileIndex,
+            Av1TransformSize transformSize,
+            Span<int> retainedCoefficients,
+            Span<Av1EncoderTransformBlockState> retainedStates,
+            out int selectedAngleDelta,
+            out Av1FilterIntraMode selectedFilterIntraMode,
+            out Av1TransformSize selectedTransformSize,
+            out long selectedCost)
+        {
+            Av1EncoderModeDecisionWorkspace<TSample> workspace =
+                this.blockWorkspace.GetModeDecisionWorkspace<TSample>();
+
+            int blockWidth = blockSize.GetWidth();
+            int blockHeight = blockSize.GetHeight();
+            int blockSampleCount = blockWidth * blockHeight;
+            int transformBlockCount = blockSampleCount / transformSize.GetSize2d();
+            Span<TSample> candidateReconstruction =
+                workspace.GetCandidateReconstruction(0)[..blockSampleCount];
+
+            Span<int> candidateCoefficients =
+                workspace.GetCandidateCoefficients(0)[..blockSampleCount];
+
+            Span<Av1EncoderTransformBlockState> candidateStates =
+                workspace.CandidateTransformBlocks[..transformBlockCount];
+
+            int contextWidth = blockSize.Get4x4WideCount();
+            int contextHeight = blockSize.Get4x4HighCount();
+            Span<byte> contexts = workspace.TransformContexts;
+            Span<byte> topContexts = contexts[..contextWidth];
+            Span<byte> leftContexts = contexts.Slice(contextWidth, contextHeight);
+            Av1NeighborArrayUnit<byte> coefficientNeighbors =
+                this.picture.LuminanceDcSignLevelCoefficientNeighbors[tileIndex];
+
+            int topIndex = coefficientNeighbors.GetTopIndex(blockOrigin);
+            int leftIndex = coefficientNeighbors.GetLeftIndex(blockOrigin);
+            int transformSizeContext = Av1TileWriter.GetTransformSizeContext(
+                this.picture.TransformFunctionContexts[tileIndex],
+                macroBlock,
+                blockOrigin,
+                blockSize);
+
+            int transformSizeRate = this.picture.Parent.FrameHeader.TransformMode == Av1TransformMode.Select
+                ? writer.GetTransformSizeCost(blockSize, transformSize, transformSizeContext)
+                : 0;
+
+            int paletteDisabledCost = Av1TileWriter.IsPaletteAllowed(
+                this.picture.Parent.FrameHeader.AllowScreenContentTools,
+                blockSize)
+                    ? writer.GetPaletteYModeCost(
+                        false,
+                        Av1TileWriter.GetPaletteBlockSizeContext(blockSize),
+                        Av1TileWriter.GetPaletteYModeContext(
+                            this.picture.PaletteContexts[tileIndex],
+                            macroBlock,
+                            blockOrigin))
+                    : 0;
+
+            int baseModeCount = LumaModeSearchOrder.Length;
+            int deltaCount = AngleDeltaSearchOrder.Length;
+            int directionalModeCount =
+                (int)Av1PredictionMode.Directional67Degrees - (int)Av1PredictionMode.Vertical + 1;
+
+            int candidateCount = this.effort switch
+            {
+                0 => 1,
+                1 => baseModeCount,
+                _ => baseModeCount + (directionalModeCount * deltaCount)
+            };
+
+            Buffer2DRegion<TSample> sourcePlane = this.source.GetPlane(Av1Plane.Y);
+            Buffer2DRegion<TSample> reconstructionPlane = this.reconstruction.GetPlane(Av1Plane.Y);
+            long bestCost = long.MaxValue;
+            Av1PredictionMode bestMode = Av1PredictionMode.DC;
+            selectedAngleDelta = 0;
+            selectedFilterIntraMode = Av1FilterIntraMode.AllFilterIntraModes;
+            selectedTransformSize = transformSize;
+
+            // Each candidate starts from the live block-edge contexts. Transform updates remain local until
+            // that candidate wins, so later modes never inherit state from an earlier trial.
+            for (int candidateIndex = 0; candidateIndex < candidateCount; candidateIndex++)
+            {
+                Av1PredictionMode mode;
+                int angleDelta;
+                if (candidateIndex < baseModeCount)
+                {
+                    mode = LumaModeSearchOrder[candidateIndex];
+                    angleDelta = 0;
+                }
+                else
+                {
+                    int adjustedIndex = candidateIndex - baseModeCount;
+                    mode = (Av1PredictionMode)(
+                        (int)Av1PredictionMode.Vertical + (adjustedIndex / deltaCount));
+
+                    angleDelta = AngleDeltaSearchOrder[adjustedIndex % deltaCount];
+                }
+
+                coefficientNeighbors.Top.Slice(topIndex, contextWidth).CopyTo(topContexts);
+                coefficientNeighbors.Left.Slice(leftIndex, contextHeight).CopyTo(leftContexts);
+                long distortion = this.GetTiledPlaneCost(
+                    writer,
+                    macroBlock,
+                    blockOrigin,
+                    blockOrigin,
+                    blockSize,
+                    blockSize,
+                    transformSize,
+                    Av1BlockSize.Block64x64,
+                    0,
+                    0,
+                    mode,
+                    mode,
+                    angleDelta,
+                    Av1Plane.Y,
+                    sourcePlane,
+                    reconstructionPlane,
+                    candidateReconstruction,
+                    candidateCoefficients,
+                    candidateStates,
+                    topContexts,
+                    leftContexts,
+                    out int coefficientRate);
+
+                int rate = Av1TileWriter.GetLumaModeCost(writer, macroBlock, blockSize, mode, angleDelta);
+                rate += transformSizeRate + coefficientRate;
+                if (mode == Av1PredictionMode.DC)
+                {
+                    rate += paletteDisabledCost;
+                    if (Av1TileWriter.IsFilterIntraAllowedBlockSize(
+                        this.picture.Sequence.SequenceHeader.EnableFilterIntra,
+                        blockSize))
+                    {
+                        rate += writer.GetFilterIntraModeCost(
+                            Av1FilterIntraMode.AllFilterIntraModes,
+                            blockSize);
+                    }
+                }
+
+                long candidateCost = Av1RateDistortion.GetCost(this.rateMultiplier, rate, distortion);
+                if (candidateCost < bestCost)
+                {
+                    CopyTiledCandidate(
+                        candidateReconstruction,
+                        candidateCoefficients,
+                        candidateStates,
+                        reconstructionPlane,
+                        blockOrigin,
+                        blockWidth,
+                        blockHeight,
+                        transformSize,
+                        retainedCoefficients,
+                        retainedStates);
+
+                    bestCost = candidateCost;
+                    bestMode = mode;
+                    selectedAngleDelta = angleDelta;
+                }
+            }
+
+            selectedCost = bestCost;
+            return bestMode;
+        }
+
         private long GetSplitLumaCandidateCost(
             Av1SymbolEncoder writer,
             Av1MacroBlockD macroBlock,
@@ -2330,10 +2545,21 @@ internal static partial class Av1IntraSuperblockEncoder
 
             if (hasBottomLeft)
             {
-                for (int row = transformHeight; row < transformHeight * 2; row++)
+                if (transformColumn > 0)
+                {
+                    for (int row = transformHeight; row < transformHeight * 2; row++)
+                    {
+                        left[row] = candidateReconstruction[
+                            ((rowOffset + row) * planeBlockWidth) + columnOffset - 1];
+                    }
+                }
+                else
                 {
-                    left[row] = reconstructionPlane
+                    for (int row = transformHeight; row < transformHeight * 2; row++)
-                        .DangerousGetRowSpan(planeBlockOrigin.Y + rowOffset + row)[planeBlockOrigin.X - 1];
+                    {
+                        left[row] = reconstructionPlane
+                            .DangerousGetRowSpan(planeBlockOrigin.Y + rowOffset + row)[planeBlockOrigin.X - 1];
+                    }
                 }
             }
             else
@@ -2406,7 +2632,10 @@ internal static partial class Av1IntraSuperblockEncoder
                 rate += paletteDisabledCost;
             }
 
-            if (mode == Av1PredictionMode.DC && this.picture.Sequence.SequenceHeader.EnableFilterIntra)
+            if (mode == Av1PredictionMode.DC &&
+                Av1TileWriter.IsFilterIntraAllowedBlockSize(
+                    this.picture.Sequence.SequenceHeader.EnableFilterIntra,
+                    blockSize))
             {
                 rate += writer.GetFilterIntraModeCost(Av1FilterIntraMode.AllFilterIntraModes, blockSize);
             }

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

@@ -1452,7 +1452,7 @@ internal partial class Av1TileWriter
     /// <param name="enableFilterIntra">A value indicating whether the sequence enables filter-intra prediction.</param>
     /// <param name="blockSize">The block size.</param>
     /// <returns><see langword="true"/> when filter-intra prediction supports the block dimensions; otherwise, <see langword="false"/>.</returns>
-    private static bool IsFilterIntraAllowedBlockSize(bool enableFilterIntra, Av1BlockSize blockSize)
+    internal static bool IsFilterIntraAllowedBlockSize(bool enableFilterIntra, Av1BlockSize blockSize)
     {
         if (!enableFilterIntra)
         {
@@ -1573,7 +1573,7 @@ internal partial class Av1TileWriter
     /// <param name="allowScreenContentTools">A value indicating whether screen-content tools are enabled.</param>
     /// <param name="blockSize">The block size.</param>
     /// <returns><see langword="true"/> when palette mode is available for the block; otherwise, <see langword="false"/>.</returns>
-    private static bool IsPaletteAllowed(bool allowScreenContentTools, Av1BlockSize blockSize)
+    internal static bool IsPaletteAllowed(bool allowScreenContentTools, Av1BlockSize blockSize)
         => allowScreenContentTools &&
             blockSize.GetWidth() <= 64 &&
             blockSize.GetHeight() <= 64 &&
@@ -1700,19 +1700,7 @@ internal partial class Av1TileWriter
         Av1NeighborArrayUnit<byte> cr_dc_sign_level_coeff_na,
         Av1NeighborArrayUnit<byte> cb_dc_sign_level_coeff_na)
     {
-        EncodeTransformCoefficientsY(
+        EncodeTransformCoefficientRegions(
-            pcs,
-            ec_ctx,
-            writer,
-            ref blk_ptr,
-            blockOrigin,
-            intraLumaDir,
-            planeBlockSize,
-            coefficientBuffer,
-            superblockIndex,
-            luma_dc_sign_level_coeff_na);
-
-        EncodeTransformCoefficientsUv(
             pcs,
             ec_ctx,
             writer,
@@ -1722,6 +1710,7 @@ internal partial class Av1TileWriter
             planeBlockSize,
             coefficientBuffer,
             superblockIndex,
+            luma_dc_sign_level_coeff_na,
             cr_dc_sign_level_coeff_na,
             cb_dc_sign_level_coeff_na);
     }
@@ -1751,14 +1740,6 @@ internal partial class Av1TileWriter
         int superblockIndex,
         Av1NeighborArrayUnit<byte> luma_dc_sign_level_coeff_na)
     {
-        ObuFrameHeader frameHeader = pcs.Parent.FrameHeader;
-        bool usesInterTransformSet = entropyCodingContext.MacroBlockModeInfo.Block.UseIntraBlockCopy;
-        Span<int> lumaCoefficients = coefficientBuffer.GetPlaneSpan(superblockIndex, Av1Plane.Y);
-        Span<Av1EncoderTransformBlockState> lumaTransformBlocks =
-            coefficientBuffer.GetTransformBlockSpan(superblockIndex, Av1Plane.Y);
-        Av1TransformSize transformSize = entropyCodingContext.MacroBlockModeInfo.Block.TransformSize;
-        int transformBlockWidth = transformSize.Get4x4WideCount();
-        int transformBlockHeight = transformSize.Get4x4HighCount();
         Av1MacroBlockD macroBlock = entropyCodingContext.MacroBlock;
         int maximumBlocksWide = plane_bsize.GetWidth();
         int maximumBlocksHigh = plane_bsize.GetHeight();
@@ -1777,66 +1758,33 @@ internal partial class Av1TileWriter
         int maximumUnitBlocksWide = Math.Min(
             Av1BlockSize.Block64x64.Get4x4WideCount(),
             maximumBlocksWide);
+
         int maximumUnitBlocksHigh = Math.Min(
             Av1BlockSize.Block64x64.Get4x4HighCount(),
             maximumBlocksHigh);
 
-        // AV1 visits residuals in bounded 64x64 regions so transform order remains stable for 128x128 blocks.
         for (int regionRow = 0; regionRow < maximumBlocksHigh; regionRow += maximumUnitBlocksHigh)
         {
-            int unitHeight = Math.Min(maximumUnitBlocksHigh + regionRow, maximumBlocksHigh);
+            int unitBottom = Math.Min(regionRow + maximumUnitBlocksHigh, maximumBlocksHigh);
             for (int regionColumn = 0; regionColumn < maximumBlocksWide; regionColumn += maximumUnitBlocksWide)
             {
-                int unitWidth = Math.Min(maximumUnitBlocksWide + regionColumn, maximumBlocksWide);
+                int unitRight = Math.Min(regionColumn + maximumUnitBlocksWide, maximumBlocksWide);
-                for (int blockRow = regionRow; blockRow < unitHeight; blockRow += transformBlockHeight)
+                EncodeTransformCoefficientRegion(
-                {
+                    pcs,
-                    for (int blockColumn = regionColumn; blockColumn < unitWidth; blockColumn += transformBlockWidth)
+                    entropyCodingContext,
-                    {
+                    writer,
-                        int transformStateIndex = entropyCodingContext.CodedAreaSuperblock /
+                    ref blk_ptr,
-                            Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
+                    blockOrigin,
-                        ref Av1EncoderTransformBlockState transformBlock = ref lumaTransformBlocks[transformStateIndex];
+                    intraLumaDir,
-                        Point transformOrigin = blockOrigin + new Size(
+                    plane_bsize,
-                            blockColumn << Av1Constants.ModeInfoSizeLog2,
+                    Av1Plane.Y,
-                            blockRow << Av1Constants.ModeInfoSizeLog2);
+                    coefficientBuffer,
-                        Span<int> coefficients = lumaCoefficients[entropyCodingContext.CodedAreaSuperblock..];
+                    superblockIndex,
-                        Av1TransformBlockContext blockContext = GetTransformBlockContexts(
+                    luma_dc_sign_level_coeff_na,
-                            Av1ComponentType.Luminance,
+                    regionRow,
-                            luma_dc_sign_level_coeff_na,
+                    regionColumn,
-                            transformOrigin,
+                    unitBottom,
-                            plane_bsize,
+                    unitRight);
-                            transformSize);
-
-                        Av1TransformType transformType = transformBlock.TransformType;
-                        ushort endOfBlock = transformBlock.EndOfBlock;
-                        if (endOfBlock == 0)
-                        {
-                            // Empty transform blocks use the canonical transform type even when mode decision retained another candidate.
-                            transformType = transformBlock.TransformType = Av1TransformType.DctDct;
-                        }
-
-                        int culLevelY = writer.WriteCoefficients(
-                            transformSize,
-                            transformType,
-                            intraLumaDir,
-                            coefficients,
-                            Av1ComponentType.Luminance,
-                            blockContext,
-                            endOfBlock,
-                            frameHeader.UseReducedTransformSet,
-                            blk_ptr.FilterIntraMode,
-                            usesInterTransformSet);
-
-                        int transformWidth = transformSize.GetWidth();
-                        int transformHeight = transformSize.GetHeight();
-                        luma_dc_sign_level_coeff_na.UnitModeWrite(
-                            (byte)culLevelY,
-                            transformOrigin,
-                            new Size(transformWidth, transformHeight),
-                            Av1NeighborArrayUnit<byte>.UnitMask.Top | Av1NeighborArrayUnit<byte>.UnitMask.Left);
-
-                        entropyCodingContext.CodedAreaSuperblock += transformWidth * transformHeight;
-                    }
-                }
             }
         }
     }
@@ -1874,26 +1822,10 @@ internal partial class Av1TileWriter
             return;
         }
 
-        ObuFrameHeader frameHeader = pcs.Parent.FrameHeader;
-        bool usesInterTransformSet = entropyCodingContext.MacroBlockModeInfo.Block.UseIntraBlockCopy;
-        Span<int> blueCoefficients = coefficientBuffer.GetPlaneSpan(superblockIndex, Av1Plane.U);
-        Span<int> redCoefficients = coefficientBuffer.GetPlaneSpan(superblockIndex, Av1Plane.V);
-        Span<Av1EncoderTransformBlockState> blueTransformBlocks =
-            coefficientBuffer.GetTransformBlockSpan(superblockIndex, Av1Plane.U);
-        Span<Av1EncoderTransformBlockState> redTransformBlocks =
-            coefficientBuffer.GetTransformBlockSpan(superblockIndex, Av1Plane.V);
-
         int subsamplingX = colorConfig.SubSamplingX ? 1 : 0;
         int subsamplingY = colorConfig.SubSamplingY ? 1 : 0;
         Av1BlockSize chromaBlockSize = plane_bsize.GetSubsampled(colorConfig.SubSamplingX, colorConfig.SubSamplingY);
         Point chromaBlockOrigin = GetChromaBlockOrigin(blockOrigin, subsamplingX, subsamplingY);
-        Av1TransformSize chromaTransformSize = frameHeader.LosslessArray[entropyCodingContext.MacroBlockModeInfo.Block.SegmentId]
-            ? Av1TransformSize.Size4x4
-            : plane_bsize.GetMaxUvTransformSize(colorConfig.SubSamplingX, colorConfig.SubSamplingY);
-        int transformBlockWidth = chromaTransformSize.Get4x4WideCount();
-        int transformBlockHeight = chromaTransformSize.Get4x4HighCount();
-        int transformWidth = chromaTransformSize.GetWidth();
-        int transformHeight = chromaTransformSize.GetHeight();
         Av1MacroBlockD macroBlock = entropyCodingContext.MacroBlock;
         int maximumBlocksWide = chromaBlockSize.GetWidth();
         int maximumBlocksHigh = chromaBlockSize.GetHeight();
@@ -1911,81 +1843,272 @@ internal partial class Av1TileWriter
         maximumBlocksHigh >>= Av1Constants.ModeInfoSizeLog2;
         Av1BlockSize maximumUnitBlockSize =
             Av1BlockSize.Block64x64.GetSubsampled(colorConfig.SubSamplingX, colorConfig.SubSamplingY);
+
         int maximumUnitBlocksWide = Math.Min(maximumUnitBlockSize.Get4x4WideCount(), maximumBlocksWide);
         int maximumUnitBlocksHigh = Math.Min(maximumUnitBlockSize.Get4x4HighCount(), maximumBlocksHigh);
 
-        int codedAreaStart = entropyCodingContext.CodedAreaSuperblockUv;
+        for (int regionRow = 0; regionRow < maximumBlocksHigh; regionRow += maximumUnitBlocksHigh)
-        int codedAreaEnd = codedAreaStart;
+        {
+            int unitBottom = Math.Min(regionRow + maximumUnitBlocksHigh, maximumBlocksHigh);
+            for (int regionColumn = 0; regionColumn < maximumBlocksWide; regionColumn += maximumUnitBlocksWide)
+            {
+                int unitRight = Math.Min(regionColumn + maximumUnitBlocksWide, maximumBlocksWide);
+                EncodeTransformCoefficientRegion(
+                    pcs,
+                    entropyCodingContext,
+                    writer,
+                    ref blk_ptr,
+                    chromaBlockOrigin,
+                    intraLumaDir,
+                    plane_bsize,
+                    Av1Plane.U,
+                    coefficientBuffer,
+                    superblockIndex,
+                    cb_dc_sign_level_coeff_na,
+                    regionRow,
+                    regionColumn,
+                    unitBottom,
+                    unitRight);
+
+                EncodeTransformCoefficientRegion(
+                    pcs,
+                    entropyCodingContext,
+                    writer,
+                    ref blk_ptr,
+                    chromaBlockOrigin,
+                    intraLumaDir,
+                    plane_bsize,
+                    Av1Plane.V,
+                    coefficientBuffer,
+                    superblockIndex,
+                    cr_dc_sign_level_coeff_na,
+                    regionRow,
+                    regionColumn,
+                    unitBottom,
+                    unitRight);
+            }
+        }
+    }
 
-        // AV1 completes every transform in one chroma plane before advancing to the other plane.
+    private static void EncodeTransformCoefficientRegions(
-        // Both planes use the same coded-area positions because their transform geometry is identical.
+        Av1PictureControlSet pcs,
-        for (int planeIndex = 0; planeIndex < 2; planeIndex++)
+        Av1EntropyCodingContext entropyCodingContext,
+        Av1SymbolEncoder writer,
+        ref Av1EncoderBlockStruct block,
+        Point blockOrigin,
+        Av1PredictionMode intraLumaMode,
+        Av1BlockSize blockSize,
+        Av1EncoderCoefficientBuffer coefficientBuffer,
+        int superblockIndex,
+        Av1NeighborArrayUnit<byte> lumaCoefficientNeighbors,
+        Av1NeighborArrayUnit<byte> redCoefficientNeighbors,
+        Av1NeighborArrayUnit<byte> blueCoefficientNeighbors)
+    {
+        Av1MacroBlockD macroBlock = entropyCodingContext.MacroBlock;
+        int maximumBlocksWide = blockSize.GetWidth();
+        int maximumBlocksHigh = blockSize.GetHeight();
+        if (macroBlock.ToRightEdge < 0)
+        {
+            maximumBlocksWide += macroBlock.ToRightEdge >> 3;
+        }
+
+        if (macroBlock.ToBottomEdge < 0)
         {
-            bool isBluePlane = planeIndex == 0;
+            maximumBlocksHigh += macroBlock.ToBottomEdge >> 3;
-            Span<int> planeCoefficients = isBluePlane ? blueCoefficients : redCoefficients;
+        }
-            Span<Av1EncoderTransformBlockState> planeTransformBlocks =
+
-                isBluePlane ? blueTransformBlocks : redTransformBlocks;
+        maximumBlocksWide >>= Av1Constants.ModeInfoSizeLog2;
+        maximumBlocksHigh >>= Av1Constants.ModeInfoSizeLog2;
+        int maximumUnitBlocksWide = Math.Min(
+            Av1BlockSize.Block64x64.Get4x4WideCount(),
+            maximumBlocksWide);
+
+        int maximumUnitBlocksHigh = Math.Min(
+            Av1BlockSize.Block64x64.Get4x4HighCount(),
+            maximumBlocksHigh);
 
-            Av1NeighborArrayUnit<byte> coefficientNeighbors =
+        ObuColorConfig colorConfig = pcs.Sequence.SequenceHeader.ColorConfig;
-                isBluePlane ? cb_dc_sign_level_coeff_na : cr_dc_sign_level_coeff_na;
+        bool hasChroma = block.HasChroma && !colorConfig.IsMonochrome;
+        int subsamplingX = colorConfig.SubSamplingX ? 1 : 0;
+        int subsamplingY = colorConfig.SubSamplingY ? 1 : 0;
+        Point chromaBlockOrigin = GetChromaBlockOrigin(blockOrigin, subsamplingX, subsamplingY);
 
-            int codedArea = codedAreaStart;
+        // Residual syntax is region-major, then plane-major. Keeping the three plane calls together
-            for (int regionRow = 0; regionRow < maximumBlocksHigh; regionRow += maximumUnitBlocksHigh)
+        // prevents a 128x128 block from emitting later luma regions before earlier chroma regions.
+        for (int regionRow = 0; regionRow < maximumBlocksHigh; regionRow += maximumUnitBlocksHigh)
+        {
+            int unitBottom = Math.Min(regionRow + maximumUnitBlocksHigh, maximumBlocksHigh);
+            for (int regionColumn = 0; regionColumn < maximumBlocksWide; regionColumn += maximumUnitBlocksWide)
             {
-                int unitHeight = Math.Min(maximumUnitBlocksHigh + regionRow, maximumBlocksHigh);
+                int unitRight = Math.Min(regionColumn + maximumUnitBlocksWide, maximumBlocksWide);
-                for (int regionColumn = 0; regionColumn < maximumBlocksWide; regionColumn += maximumUnitBlocksWide)
+                EncodeTransformCoefficientRegion(
+                    pcs,
+                    entropyCodingContext,
+                    writer,
+                    ref block,
+                    blockOrigin,
+                    intraLumaMode,
+                    blockSize,
+                    Av1Plane.Y,
+                    coefficientBuffer,
+                    superblockIndex,
+                    lumaCoefficientNeighbors,
+                    regionRow,
+                    regionColumn,
+                    unitBottom,
+                    unitRight);
+
+                if (hasChroma)
                 {
-                    int unitWidth = Math.Min(maximumUnitBlocksWide + regionColumn, maximumBlocksWide);
+                    int chromaRegionRow = regionRow >> subsamplingY;
-                    for (int blockRow = regionRow; blockRow < unitHeight; blockRow += transformBlockHeight)
+                    int chromaRegionColumn = regionColumn >> subsamplingX;
-                    {
+                    int chromaUnitBottom = unitBottom >> subsamplingY;
-                        for (int blockColumn = regionColumn; blockColumn < unitWidth; blockColumn += transformBlockWidth)
+                    int chromaUnitRight = unitRight >> subsamplingX;
-                        {
+                    EncodeTransformCoefficientRegion(
-                            int transformStateIndex =
+                        pcs,
-                                codedArea / Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
+                        entropyCodingContext,
-
+                        writer,
-                            ref Av1EncoderTransformBlockState transformBlock =
+                        ref block,
-                                ref planeTransformBlocks[transformStateIndex];
+                        chromaBlockOrigin,
-
+                        intraLumaMode,
-                            Point chromaOrigin = chromaBlockOrigin + new Size(
+                        blockSize,
-                                blockColumn << Av1Constants.ModeInfoSizeLog2,
+                        Av1Plane.U,
-                                blockRow << Av1Constants.ModeInfoSizeLog2);
+                        coefficientBuffer,
-
+                        superblockIndex,
-                            Span<int> coefficients = planeCoefficients[codedArea..];
+                        blueCoefficientNeighbors,
-                            Av1TransformBlockContext blockContext = GetTransformBlockContexts(
+                        chromaRegionRow,
-                                Av1ComponentType.Chroma,
+                        chromaRegionColumn,
-                                coefficientNeighbors,
+                        chromaUnitBottom,
-                                chromaOrigin,
+                        chromaUnitRight);
-                                chromaBlockSize,
+
-                                chromaTransformSize);
+                    EncodeTransformCoefficientRegion(
-
+                        pcs,
-                            int culLevel = writer.WriteCoefficients(
+                        entropyCodingContext,
-                                chromaTransformSize,
+                        writer,
-                                transformBlock.TransformType,
+                        ref block,
-                                intraLumaDir,
+                        chromaBlockOrigin,
-                                coefficients,
+                        intraLumaMode,
-                                Av1ComponentType.Chroma,
+                        blockSize,
-                                blockContext,
+                        Av1Plane.V,
-                                transformBlock.EndOfBlock,
+                        coefficientBuffer,
-                                frameHeader.UseReducedTransformSet,
+                        superblockIndex,
-                                blk_ptr.FilterIntraMode,
+                        redCoefficientNeighbors,
-                                usesInterTransformSet);
+                        chromaRegionRow,
-
+                        chromaRegionColumn,
-                            coefficientNeighbors.UnitModeWrite(
+                        chromaUnitBottom,
-                                (byte)culLevel,
+                        chromaUnitRight);
-                                chromaOrigin,
-                                new Size(transformWidth, transformHeight),
-                                Av1NeighborArrayUnit<byte>.UnitMask.Top | Av1NeighborArrayUnit<byte>.UnitMask.Left);
-
-                            codedArea += transformWidth * transformHeight;
-                        }
-                    }
                 }
             }
+        }
+    }
+
+    private static void EncodeTransformCoefficientRegion(
+        Av1PictureControlSet pcs,
+        Av1EntropyCodingContext entropyCodingContext,
+        Av1SymbolEncoder writer,
+        ref Av1EncoderBlockStruct block,
+        Point planeBlockOrigin,
+        Av1PredictionMode intraLumaMode,
+        Av1BlockSize lumaBlockSize,
+        Av1Plane plane,
+        Av1EncoderCoefficientBuffer coefficientBuffer,
+        int superblockIndex,
+        Av1NeighborArrayUnit<byte> coefficientNeighbors,
+        int regionRow,
+        int regionColumn,
+        int unitBottom,
+        int unitRight)
+    {
+        ObuFrameHeader frameHeader = pcs.Parent.FrameHeader;
+        ObuColorConfig colorConfig = pcs.Sequence.SequenceHeader.ColorConfig;
+        bool isLuma = plane == Av1Plane.Y;
+        Av1BlockSize planeBlockSize = isLuma
+            ? lumaBlockSize
+            : lumaBlockSize.GetSubsampled(colorConfig.SubSamplingX, colorConfig.SubSamplingY);
 
-            codedAreaEnd = codedArea;
+        Av1TransformSize transformSize = isLuma
+            ? entropyCodingContext.MacroBlockModeInfo.Block.TransformSize
+            : frameHeader.LosslessArray[entropyCodingContext.MacroBlockModeInfo.Block.SegmentId]
+                ? Av1TransformSize.Size4x4
+                : lumaBlockSize.GetMaxUvTransformSize(colorConfig.SubSamplingX, colorConfig.SubSamplingY);
+
+        int transformBlockWidth = transformSize.Get4x4WideCount();
+        int transformBlockHeight = transformSize.Get4x4HighCount();
+        int transformWidth = transformSize.GetWidth();
+        int transformHeight = transformSize.GetHeight();
+        bool usesInterTransformSet = entropyCodingContext.MacroBlockModeInfo.Block.UseIntraBlockCopy;
+        Av1ComponentType componentType = isLuma
+            ? Av1ComponentType.Luminance
+            : Av1ComponentType.Chroma;
+
+        Span<int> planeCoefficients = coefficientBuffer.GetPlaneSpan(superblockIndex, plane);
+        Span<Av1EncoderTransformBlockState> planeTransformBlocks =
+            coefficientBuffer.GetTransformBlockSpan(superblockIndex, plane);
+
+        int codedArea = isLuma
+            ? entropyCodingContext.CodedAreaSuperblock
+            : entropyCodingContext.CodedAreaSuperblockUv;
+
+        for (int blockRow = regionRow; blockRow < unitBottom; blockRow += transformBlockHeight)
+        {
+            for (int blockColumn = regionColumn; blockColumn < unitRight; blockColumn += transformBlockWidth)
+            {
+                int transformStateIndex =
+                    codedArea / Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
+
+                ref Av1EncoderTransformBlockState transformBlock =
+                    ref planeTransformBlocks[transformStateIndex];
+
+                Point transformOrigin = planeBlockOrigin + new Size(
+                    blockColumn << Av1Constants.ModeInfoSizeLog2,
+                    blockRow << Av1Constants.ModeInfoSizeLog2);
+
+                Span<int> coefficients = planeCoefficients[codedArea..];
+                Av1TransformBlockContext blockContext = GetTransformBlockContexts(
+                    componentType,
+                    coefficientNeighbors,
+                    transformOrigin,
+                    planeBlockSize,
+                    transformSize);
+
+                Av1TransformType transformType = transformBlock.TransformType;
+                if (isLuma && transformBlock.EndOfBlock == 0)
+                {
+                    // Empty luma transforms carry no transform-type symbol, so retain the canonical state.
+                    transformType = transformBlock.TransformType = Av1TransformType.DctDct;
+                }
+
+                int culLevel = writer.WriteCoefficients(
+                    transformSize,
+                    transformType,
+                    intraLumaMode,
+                    coefficients,
+                    componentType,
+                    blockContext,
+                    transformBlock.EndOfBlock,
+                    frameHeader.UseReducedTransformSet,
+                    block.FilterIntraMode,
+                    usesInterTransformSet);
+
+                coefficientNeighbors.UnitModeWrite(
+                    (byte)culLevel,
+                    transformOrigin,
+                    new Size(transformWidth, transformHeight),
+                    Av1NeighborArrayUnit<byte>.UnitMask.Top | Av1NeighborArrayUnit<byte>.UnitMask.Left);
+
+                codedArea += transformWidth * transformHeight;
+            }
         }
 
-        entropyCodingContext.CodedAreaSuperblockUv = codedAreaEnd;
+        if (isLuma)
+        {
+            entropyCodingContext.CodedAreaSuperblock = codedArea;
+        }
+        else if (plane == Av1Plane.V)
+        {
+            // U and V share the same per-plane coded-area positions; advance only after V completes the region.
+            entropyCodingContext.CodedAreaSuperblockUv = codedArea;
+        }
     }
 
     /// <summary>

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

@@ -357,6 +357,80 @@ public class Av1EncoderFrameTests
         Assert.Equal(new Size(Size, Size), decoded.Size);
     }
 
+    [Theory]
+    [InlineData(Yuv400)]
+    [InlineData(Yuv444)]
+    public void EncodeEffortTenSelectsOneHundredTwentyEightByOneHundredTwentyEightBlock(int colorFormatValue)
+    {
+        const int Size = 128;
+        Av1ColorFormat colorFormat = (Av1ColorFormat)colorFormatValue;
+        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++)
+            {
+                row[x] = new Rgba32(180, 64, 220);
+            }
+        }
+
+        using MemoryStream stream = new();
+        ObuSequenceHeader sequenceHeader = Av1FrameEncoder.Encode(
+            Configuration.Default,
+            source.Frames.RootFrame,
+            stream,
+            CreateColorConfig(Av1BitDepth.EightBit, colorFormat),
+            qIndex: 4,
+            effort: 10);
+
+        Assert.True(sequenceHeader.Use128x128Superblock);
+        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 = 0; modeInfoY < 32; modeInfoY++)
+        {
+            for (int modeInfoX = 0; modeInfoX < 32; modeInfoX++)
+            {
+                Assert.Equal(
+                    Av1BlockSize.Block128x128,
+                    frameInfo.GetModeInfoAt(new Point(modeInfoX, modeInfoY)).BlockSize);
+            }
+        }
+
+        Assert.Equal(new Size(Size, Size), decoded.Size);
+    }
+
+    [Fact]
+    public void EncodeEffortTenSearchesHighBitDepthOneHundredTwentyEightRoot()
+    {
+        const int Size = 128;
+        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++)
+            {
+                row[x] = new Rgba32(180, 64, 220);
+            }
+        }
+
+        using MemoryStream stream = new();
+        ObuSequenceHeader sequenceHeader = Av1FrameEncoder.Encode(
+            Configuration.Default,
+            source.Frames.RootFrame,
+            stream,
+            CreateColorConfig(Av1BitDepth.TwelveBit, Av1ColorFormat.Yuv444),
+            qIndex: 4,
+            effort: 10);
+
+        Assert.True(sequenceHeader.Use128x128Superblock);
+        byte[] payload = stream.ToArray();
+        using Av1Decoder decoder = new(Configuration.Default);
+        using Image<Rgba32> decoded = decoder.Decode<Rgba32>(payload);
+        Assert.Equal(new Size(Size, Size), decoded.Size);
+    }
+
     [Theory]
     [InlineData(EightBit)]
     [InlineData(TenBit)]

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

@@ -618,19 +618,25 @@ public class Av1TransformBlockEncoderTests
                 workspace.GetIntraBlockCopyWorkspace<ushort>();
 
             Assert.Equal(
-                (2 * Av1EncoderModeDecisionWorkspace<ushort>.MaximumBlockDimension) + 1,
+                (2 * Av1Constants.MaxTransformSize) + 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(Av1EncoderModeDecisionWorkspace<ushort>.MaximumTransformSampleCount, modeWorkspace.Prediction.Length);
+            Assert.Equal(Av1EncoderModeDecisionWorkspace<ushort>.MaximumTransformSampleCount, modeWorkspace.Residual.Length);
+            Assert.Equal(Av1EncoderModeDecisionWorkspace<ushort>.MaximumCandidateTransformBlockCount, modeWorkspace.CandidateTransformBlocks.Length);
 
             // CfL is unavailable above 32x32, so its scratch remains fixed while larger partitions are enabled.
             Assert.Equal(Av1ChromaFromLumaContext.BufferLength, 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);
+            int maximumPaletteSampleCount =
-            Assert.Equal(Av1EncoderModeDecisionWorkspace<ushort>.MaximumSampleCount, paletteWorkspace.AlternateIndices.Length);
+                Av1BlockSize.Block64x64.GetWidth() * Av1BlockSize.Block64x64.GetHeight();
+
+            Assert.Equal(maximumPaletteSampleCount, paletteWorkspace.GetPrediction(1).Length);
+            Assert.Equal(maximumPaletteSampleCount, paletteWorkspace.AlternateIndices.Length);
 
             // Conventional mode search and IBC are sequential, so their typed views intentionally alias one owner region.
             modeWorkspace.GetReferenceSamples(0)[0] = 123;