Extend AV1 partition search through 64x64

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

@@ -18,7 +18,7 @@ internal readonly ref struct Av1EncoderModeDecisionWorkspace<TSample>
     /// <summary>
     /// The largest coding-block dimension evaluated directly by the current partition search.
     /// </summary>
-    public const int MaximumBlockDimension = 32;
+    public const int MaximumBlockDimension = 64;
 
     /// <summary>
     /// The maximum number of samples in one directly evaluated coding block.
@@ -30,6 +30,11 @@ internal readonly ref struct Av1EncoderModeDecisionWorkspace<TSample>
     /// </summary>
     public const int CandidateTransformBlockCount = 4;
 
+    /// <summary>
+    /// The maximum number of transform states needed while evaluating both chroma planes of one 64x64 block.
+    /// </summary>
+    public const int MaximumCandidateTransformBlockCount = 8;
+
     /// <summary>
     /// The required workspace length in signed-integer storage elements.
     /// </summary>
@@ -43,10 +48,10 @@ internal readonly ref struct Av1EncoderModeDecisionWorkspace<TSample>
     private const int CandidateCoefficientStorageOffset = CandidateSampleStorageOffset + CandidateSampleStorageLength;
     private const int CandidateCoefficientStorageLength = 2 * MaximumSampleCount;
     private const int CandidateTransformBlockStorageOffset = CandidateCoefficientStorageOffset + CandidateCoefficientStorageLength;
-    private const int CandidateTransformBlockStorageLength = CandidateTransformBlockCount;
+    private const int CandidateTransformBlockStorageLength = MaximumCandidateTransformBlockCount;
     private const int TransformContextStorageOffset = CandidateTransformBlockStorageOffset + CandidateTransformBlockStorageLength;
     private const int TransformContextStorageLength =
-        2 * (MaximumBlockDimension >> Av1Constants.ModeInfoSizeLog2) * sizeof(byte) / sizeof(int);
+        4 * (MaximumBlockDimension >> Av1Constants.ModeInfoSizeLog2) * sizeof(byte) / sizeof(int);
 
     private const int TransientStorageOffset = TransformContextStorageOffset + TransformContextStorageLength;
     private const int ChromaFromLumaSampleCount = Av1ChromaFromLumaContext.BufferLength;
@@ -100,14 +105,14 @@ internal readonly ref struct Av1EncoderModeDecisionWorkspace<TSample>
         => new(this.storage[TransientStorageOffset..]);
 
     /// <summary>
-    /// Gets the transform state retained while evaluating a uniform 4x4 luma layout.
+    /// Gets the transform states retained while evaluating a multi-transform candidate.
     /// </summary>
     public Span<Av1EncoderTransformBlockState> CandidateTransformBlocks
         => MemoryMarshal.Cast<int, Av1EncoderTransformBlockState>(
             this.storage.Slice(CandidateTransformBlockStorageOffset, CandidateTransformBlockStorageLength));
 
     /// <summary>
-    /// Gets the two above and two left coefficient contexts used by a uniform 4x4 luma layout.
+    /// Gets coefficient-context edges shared by mutually exclusive luma and chroma transform trials.
     /// </summary>
     public Span<byte> TransformContexts
         => MemoryMarshal.AsBytes(this.storage.Slice(TransformContextStorageOffset, TransformContextStorageLength));

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

@@ -52,8 +52,8 @@ internal static partial class Av1IntraSuperblockEncoder
             Av1TransformSize transformSize,
             Span<int> retainedBlueCoefficients,
             Span<int> retainedRedCoefficients,
-            ref Av1EncoderTransformBlockState retainedBlueState,
+            Span<Av1EncoderTransformBlockState> retainedBlueStates,
-            ref Av1EncoderTransformBlockState retainedRedState,
+            Span<Av1EncoderTransformBlockState> retainedRedStates,
             ref Av1EncoderPaletteInfo paletteInfo,
             out int selectedAngleDelta,
             out byte selectedChromaFromLumaIndex,
@@ -69,6 +69,34 @@ internal static partial class Av1IntraSuperblockEncoder
             int width = transformSize.GetWidth();
             int height = transformSize.GetHeight();
             int sampleCount = transformSize.GetSize2d();
+            Av1BlockSize chromaBlockSize = blockSize.GetSubsampled(
+                colorConfig.SubSamplingX,
+                colorConfig.SubSamplingY);
+
+            if (chromaBlockSize.GetWidth() > width || chromaBlockSize.GetHeight() > height)
+            {
+                return this.SelectTiledChromaMode(
+                    writer,
+                    macroBlock,
+                    modeInfo,
+                    lumaOrigin,
+                    chromaOrigin,
+                    blockSize,
+                    chromaBlockSize,
+                    tileIndex,
+                    lumaMode,
+                    transformSize,
+                    retainedBlueCoefficients,
+                    retainedRedCoefficients,
+                    retainedBlueStates,
+                    retainedRedStates,
+                    paletteInfo,
+                    out selectedAngleDelta,
+                    out selectedChromaFromLumaIndex,
+                    out selectedChromaFromLumaSigns,
+                    out selectedCost);
+            }
+
             int modeInfoRow = lumaOrigin.Y >> Av1Constants.ModeInfoSizeLog2;
             int modeInfoColumn = lumaOrigin.X >> Av1Constants.ModeInfoSizeLog2;
             bool hasLeft = macroBlock.IsLeftAvailable;
@@ -152,7 +180,6 @@ 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);
             Av1TransformBlockContext blueContext = Av1TileWriter.GetTransformBlockContexts(
                 Av1ComponentType.Chroma,
                 this.picture.CbDcSignLevelCoefficientNeighbors[tileIndex],
@@ -255,7 +282,7 @@ internal static partial class Av1IntraSuperblockEncoder
                         retainedBlueCoefficients,
                         transformSize,
                         candidateBlueState,
-                        ref retainedBlueState);
+                        ref retainedBlueStates[0]);
 
                     CopyCandidate(
                         candidateRedReconstruction,
@@ -265,7 +292,7 @@ internal static partial class Av1IntraSuperblockEncoder
                         retainedRedCoefficients,
                         transformSize,
                         candidateRedState,
-                        ref retainedRedState);
+                        ref retainedRedStates[0]);
 
                     bestCost = candidateCost;
                     bestMode = chromaMode;
@@ -447,7 +474,7 @@ internal static partial class Av1IntraSuperblockEncoder
                         retainedBlueCoefficients,
                         transformSize,
                         candidateBlueState,
-                        ref retainedBlueState);
+                        ref retainedBlueStates[0]);
 
                     Av1EncoderTransformBlockState candidateRedState = default;
                     _ = this.GetChromaFromLumaPlaneCost(
@@ -474,7 +501,7 @@ internal static partial class Av1IntraSuperblockEncoder
                         retainedRedCoefficients,
                         transformSize,
                         candidateRedState,
-                        ref retainedRedState);
+                        ref retainedRedStates[0]);
                 }
             }
 
@@ -498,8 +525,8 @@ internal static partial class Av1IntraSuperblockEncoder
                     candidateRedCoefficients[..sampleCount],
                     retainedBlueCoefficients,
                     retainedRedCoefficients,
-                    ref retainedBlueState,
+                    ref retainedBlueStates[0],
-                    ref retainedRedState,
+                    ref retainedRedStates[0],
                     ref bestCost,
                     ref paletteInfo))
             {
@@ -513,6 +540,408 @@ internal static partial class Av1IntraSuperblockEncoder
             return bestMode;
         }
 
+        private Av1ChromaPredictionMode SelectTiledChromaMode(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Av1MacroBlockModeInfo modeInfo,
+            Point lumaOrigin,
+            Point chromaOrigin,
+            Av1BlockSize blockSize,
+            Av1BlockSize chromaBlockSize,
+            ushort tileIndex,
+            Av1PredictionMode lumaMode,
+            Av1TransformSize transformSize,
+            Span<int> retainedBlueCoefficients,
+            Span<int> retainedRedCoefficients,
+            Span<Av1EncoderTransformBlockState> retainedBlueStates,
+            Span<Av1EncoderTransformBlockState> retainedRedStates,
+            Av1EncoderPaletteInfo paletteInfo,
+            out int selectedAngleDelta,
+            out byte selectedChromaFromLumaIndex,
+            out sbyte selectedChromaFromLumaSigns,
+            out long selectedCost)
+        {
+            Av1EncoderModeDecisionWorkspace<TSample> workspace =
+                this.blockWorkspace.GetModeDecisionWorkspace<TSample>();
+
+            ObuColorConfig colorConfig = this.picture.Sequence.SequenceHeader.ColorConfig;
+            int subsamplingX = colorConfig.SubSamplingX ? 1 : 0;
+            int subsamplingY = colorConfig.SubSamplingY ? 1 : 0;
+            int blockWidth = chromaBlockSize.GetWidth();
+            int blockHeight = chromaBlockSize.GetHeight();
+            int blockSampleCount = blockWidth * blockHeight;
+            int transformWidth = transformSize.GetWidth();
+            int transformHeight = transformSize.GetHeight();
+            int transformColumnCount = blockWidth / transformWidth;
+            int transformRowCount = blockHeight / transformHeight;
+            int transformBlockCount = transformColumnCount * transformRowCount;
+            Span<TSample> candidateBlueReconstruction =
+                workspace.GetCandidateReconstruction(0)[..blockSampleCount];
+
+            Span<TSample> candidateRedReconstruction =
+                workspace.GetCandidateReconstruction(1)[..blockSampleCount];
+
+            Span<int> candidateBlueCoefficients =
+                workspace.GetCandidateCoefficients(0)[..blockSampleCount];
+
+            Span<int> candidateRedCoefficients =
+                workspace.GetCandidateCoefficients(1)[..blockSampleCount];
+
+            Span<Av1EncoderTransformBlockState> candidateStates = workspace.CandidateTransformBlocks;
+            Span<Av1EncoderTransformBlockState> candidateBlueStates =
+                candidateStates[..transformBlockCount];
+
+            Span<Av1EncoderTransformBlockState> candidateRedStates =
+                candidateStates.Slice(transformBlockCount, transformBlockCount);
+
+            int contextWidth = chromaBlockSize.Get4x4WideCount();
+            int contextHeight = chromaBlockSize.Get4x4HighCount();
+            Span<byte> contexts = workspace.TransformContexts;
+            Span<byte> blueTopContexts = contexts[..contextWidth];
+            Span<byte> blueLeftContexts = contexts.Slice(contextWidth, contextHeight);
+            Span<byte> redTopContexts = contexts.Slice(contextWidth + contextHeight, contextWidth);
+            Span<byte> redLeftContexts = contexts.Slice(
+                (2 * contextWidth) + contextHeight,
+                contextHeight);
+
+            Av1NeighborArrayUnit<byte> blueNeighbors =
+                this.picture.CbDcSignLevelCoefficientNeighbors[tileIndex];
+
+            Av1NeighborArrayUnit<byte> redNeighbors =
+                this.picture.CrDcSignLevelCoefficientNeighbors[tileIndex];
+
+            int blueTopIndex = blueNeighbors.GetTopIndex(chromaOrigin);
+            int blueLeftIndex = blueNeighbors.GetLeftIndex(chromaOrigin);
+            int redTopIndex = redNeighbors.GetTopIndex(chromaOrigin);
+            int redLeftIndex = redNeighbors.GetLeftIndex(chromaOrigin);
+            Buffer2DRegion<TSample> blueSource = this.source.GetPlane(Av1Plane.U);
+            Buffer2DRegion<TSample> redSource = this.source.GetPlane(Av1Plane.V);
+            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
+                ? writer.GetPaletteUvModeCost(false, hasLumaPalette)
+                : 0;
+
+            int baseModeCount = ChromaModeSearchOrder.Length;
+            int deltaCount = AngleDeltaSearchOrder.Length;
+            int directionalModeCount =
+                (int)Av1ChromaPredictionMode.Directional67Degrees -
+                (int)Av1ChromaPredictionMode.Vertical +
+                1;
+
+            int candidateCount = baseModeCount + (directionalModeCount * deltaCount);
+            long bestCost = long.MaxValue;
+            Av1ChromaPredictionMode bestMode = Av1ChromaPredictionMode.DC;
+            selectedAngleDelta = 0;
+            selectedChromaFromLumaIndex = 0;
+            selectedChromaFromLumaSigns = 0;
+
+            // A large chroma block is predicted and transformed in the same raster order used by the tile
+            // writer. Each completed transform supplies both reconstructed edges and coefficient contexts.
+            for (int candidateIndex = 0; candidateIndex < candidateCount; candidateIndex++)
+            {
+                Av1ChromaPredictionMode chromaMode;
+                int angleDelta;
+                if (candidateIndex < baseModeCount)
+                {
+                    chromaMode = ChromaModeSearchOrder[candidateIndex];
+                    angleDelta = 0;
+                }
+                else
+                {
+                    int adjustedIndex = candidateIndex - baseModeCount;
+                    chromaMode = (Av1ChromaPredictionMode)(
+                        (int)Av1ChromaPredictionMode.Vertical + (adjustedIndex / deltaCount));
+
+                    angleDelta = AngleDeltaSearchOrder[adjustedIndex % deltaCount];
+                }
+
+                blueNeighbors.Top.Slice(blueTopIndex, contextWidth).CopyTo(blueTopContexts);
+                blueNeighbors.Left.Slice(blueLeftIndex, contextHeight).CopyTo(blueLeftContexts);
+                redNeighbors.Top.Slice(redTopIndex, contextWidth).CopyTo(redTopContexts);
+                redNeighbors.Left.Slice(redLeftIndex, contextHeight).CopyTo(redLeftContexts);
+                Av1PredictionMode predictionMode = chromaMode.ToLumaMode();
+                long distortion = this.GetTiledChromaPlaneCost(
+                    writer,
+                    macroBlock,
+                    lumaOrigin,
+                    chromaOrigin,
+                    blockSize,
+                    chromaBlockSize,
+                    transformSize,
+                    subsamplingX,
+                    subsamplingY,
+                    lumaMode,
+                    predictionMode,
+                    angleDelta,
+                    Av1Plane.U,
+                    blueSource,
+                    blueReconstruction,
+                    candidateBlueReconstruction,
+                    candidateBlueCoefficients,
+                    candidateBlueStates,
+                    blueTopContexts,
+                    blueLeftContexts,
+                    out int blueRate);
+
+                distortion += this.GetTiledChromaPlaneCost(
+                    writer,
+                    macroBlock,
+                    lumaOrigin,
+                    chromaOrigin,
+                    blockSize,
+                    chromaBlockSize,
+                    transformSize,
+                    subsamplingX,
+                    subsamplingY,
+                    lumaMode,
+                    predictionMode,
+                    angleDelta,
+                    Av1Plane.V,
+                    redSource,
+                    redReconstruction,
+                    candidateRedReconstruction,
+                    candidateRedCoefficients,
+                    candidateRedStates,
+                    redTopContexts,
+                    redLeftContexts,
+                    out int redRate);
+
+                int rate = Av1TileWriter.GetChromaModeCost(
+                    writer,
+                    this.picture.Parent.FrameHeader,
+                    colorConfig,
+                    modeInfo,
+                    blockSize,
+                    lumaMode,
+                    chromaMode,
+                    angleDelta);
+
+                rate += blueRate + redRate;
+                if (chromaMode == Av1ChromaPredictionMode.DC)
+                {
+                    rate += paletteDisabledCost;
+                }
+
+                long candidateCost = Av1RateDistortion.GetCost(this.rateMultiplier, rate, distortion);
+                if (candidateCost < bestCost)
+                {
+                    CopyTiledCandidate(
+                        candidateBlueReconstruction,
+                        candidateBlueCoefficients,
+                        candidateBlueStates,
+                        blueReconstruction,
+                        chromaOrigin,
+                        blockWidth,
+                        blockHeight,
+                        transformSize,
+                        retainedBlueCoefficients,
+                        retainedBlueStates);
+
+                    CopyTiledCandidate(
+                        candidateRedReconstruction,
+                        candidateRedCoefficients,
+                        candidateRedStates,
+                        redReconstruction,
+                        chromaOrigin,
+                        blockWidth,
+                        blockHeight,
+                        transformSize,
+                        retainedRedCoefficients,
+                        retainedRedStates);
+
+                    bestCost = candidateCost;
+                    bestMode = chromaMode;
+                    selectedAngleDelta = angleDelta;
+                }
+            }
+
+            selectedCost = bestCost;
+            return bestMode;
+        }
+
+        private long GetTiledChromaPlaneCost(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Point lumaOrigin,
+            Point chromaOrigin,
+            Av1BlockSize blockSize,
+            Av1BlockSize chromaBlockSize,
+            Av1TransformSize transformSize,
+            int subsamplingX,
+            int subsamplingY,
+            Av1PredictionMode lumaMode,
+            Av1PredictionMode predictionMode,
+            int angleDelta,
+            Av1Plane plane,
+            Buffer2DRegion<TSample> source,
+            Buffer2DRegion<TSample> reconstruction,
+            Span<TSample> candidateReconstruction,
+            Span<int> candidateCoefficients,
+            Span<Av1EncoderTransformBlockState> candidateStates,
+            Span<byte> topContexts,
+            Span<byte> leftContexts,
+            out int rate)
+        {
+            Av1EncoderModeDecisionWorkspace<TSample> workspace =
+                this.blockWorkspace.GetModeDecisionWorkspace<TSample>();
+
+            int blockWidth = chromaBlockSize.GetWidth();
+            int blockHeight = chromaBlockSize.GetHeight();
+            int transformWidth = transformSize.GetWidth();
+            int transformHeight = transformSize.GetHeight();
+            int transformSampleCount = transformSize.GetSize2d();
+            int transformWidth4x4 = transformSize.Get4x4WideCount();
+            int transformHeight4x4 = transformSize.Get4x4HighCount();
+            Av1TransformType transformType = Av1SymbolContextHelper.GetDefaultIntraTransformType(
+                predictionMode,
+                transformSize,
+                this.picture.Parent.FrameHeader.UseReducedTransformSet);
+
+            Av1ComponentType componentType = Av1ComponentType.Chroma;
+            Span<TSample> prediction = workspace.Prediction[..transformSampleCount];
+            Span<short> residual = workspace.Residual[..transformSampleCount];
+            Span<TSample> aboveStorage = workspace.GetReferenceSamples(0);
+            Span<TSample> leftStorage = workspace.GetReferenceSamples(1);
+            int coefficientOffset = 0;
+            int transformIndex = 0;
+            long distortion = 0;
+            rate = 0;
+            for (int transformRow = 0; transformRow < blockHeight / transformHeight; transformRow++)
+            {
+                int rowOffset = transformRow * transformHeight;
+                for (int transformColumn = 0; transformColumn < blockWidth / transformWidth; transformColumn++)
+                {
+                    int columnOffset = transformColumn * transformWidth;
+                    int reconstructionOffset = (rowOffset * blockWidth) + columnOffset;
+                    Point transformOrigin = chromaOrigin + new Size(columnOffset, rowOffset);
+                    this.PrepareTransformReferenceSamples(
+                        reconstruction,
+                        lumaOrigin,
+                        chromaOrigin,
+                        blockSize,
+                        macroBlock,
+                        transformRow,
+                        transformColumn,
+                        blockWidth,
+                        transformSize,
+                        subsamplingX,
+                        subsamplingY,
+                        candidateReconstruction,
+                        aboveStorage,
+                        leftStorage,
+                        out bool hasLeft,
+                        out bool hasAbove);
+
+                    TOperator.PrepareIntra(
+                        this.blockWorkspace,
+                        source,
+                        transformOrigin,
+                        prediction,
+                        aboveStorage.Slice(1, transformWidth * 2),
+                        leftStorage.Slice(1, transformHeight * 2),
+                        hasLeft,
+                        hasAbove,
+                        predictionMode,
+                        angleDelta,
+                        residual,
+                        transformSize,
+                        this.bitDepth);
+
+                    Av1TransformBlockContext blockContext = Av1TileWriter.GetTransformBlockContexts(
+                        componentType,
+                        topContexts.Slice(transformColumn * transformWidth4x4, transformWidth4x4),
+                        leftContexts.Slice(transformRow * transformHeight4x4, transformHeight4x4),
+                        chromaBlockSize,
+                        transformSize);
+
+                    Span<int> transformCoefficients = candidateCoefficients.Slice(
+                        coefficientOffset,
+                        transformSampleCount);
+
+                    ref Av1EncoderTransformBlockState state = ref candidateStates[transformIndex++];
+                    distortion += TOperator.EncodePredictionCandidate(
+                        this.blockWorkspace,
+                        source,
+                        transformOrigin,
+                        prediction,
+                        residual,
+                        candidateReconstruction[reconstructionOffset..],
+                        blockWidth,
+                        transformCoefficients,
+                        transformSize,
+                        transformType,
+                        plane,
+                        this.quantization.QIndex[0],
+                        this.quantization.DeltaQDc[(int)plane],
+                        this.quantization.DeltaQAc[(int)plane],
+                        this.bitDepth,
+                        ref state);
+
+                    rate += writer.GetCoefficientCost(
+                        transformSize,
+                        transformType,
+                        lumaMode,
+                        transformCoefficients,
+                        componentType,
+                        blockContext,
+                        state.EndOfBlock,
+                        this.picture.Parent.FrameHeader.UseReducedTransformSet,
+                        Av1FilterIntraMode.AllFilterIntraModes,
+                        usesInterTransformSet: false);
+
+                    byte coefficientContext = Av1SymbolContextHelper.GetCoefficientContext(
+                        transformCoefficients,
+                        transformSize,
+                        transformType,
+                        state.EndOfBlock);
+
+                    topContexts
+                        .Slice(transformColumn * transformWidth4x4, transformWidth4x4)
+                        .Fill(coefficientContext);
+
+                    leftContexts
+                        .Slice(transformRow * transformHeight4x4, transformHeight4x4)
+                        .Fill(coefficientContext);
+
+                    coefficientOffset += transformSampleCount;
+                }
+            }
+
+            return distortion;
+        }
+
+        private static void CopyTiledCandidate(
+            ReadOnlySpan<TSample> candidateReconstruction,
+            ReadOnlySpan<int> candidateCoefficients,
+            ReadOnlySpan<Av1EncoderTransformBlockState> candidateStates,
+            Buffer2DRegion<TSample> reconstruction,
+            Point blockOrigin,
+            int blockWidth,
+            int blockHeight,
+            Av1TransformSize transformSize,
+            Span<int> retainedCoefficients,
+            Span<Av1EncoderTransformBlockState> retainedStates)
+        {
+            int blockSampleCount = blockWidth * blockHeight;
+            int transformStateStride =
+                transformSize.GetSize2d() /
+                Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
+
+            candidateCoefficients[..blockSampleCount].CopyTo(retainedCoefficients);
+            for (int transformIndex = 0; transformIndex < candidateStates.Length; transformIndex++)
+            {
+                retainedStates[transformIndex * transformStateStride] = candidateStates[transformIndex];
+            }
+
+            for (int row = 0; row < blockHeight; row++)
+            {
+                candidateReconstruction.Slice(row * blockWidth, blockWidth)
+                    .CopyTo(reconstruction.DangerousGetRowSpan(blockOrigin.Y + row).Slice(blockOrigin.X, blockWidth));
+            }
+        }
+
         private long GetChromaFromLumaPlaneCost(
             Av1SymbolEncoder writer,
             Av1PredictionMode lumaMode,

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

@@ -184,7 +184,7 @@ internal static partial class Av1IntraSuperblockEncoder
             Av1PartitionType preparedPartition)
         {
             bool searchPartition = blockSize is Av1BlockSize.Block8x8 or Av1BlockSize.Block16x16 ||
-                (this.effort == 10 && blockSize == Av1BlockSize.Block32x32);
+                (this.effort == 10 && blockSize is Av1BlockSize.Block32x32 or Av1BlockSize.Block64x64);
 
             if (!searchPartition)
             {
@@ -652,6 +652,10 @@ internal static partial class Av1IntraSuperblockEncoder
                 colorConfig.SubSamplingX,
                 colorConfig.SubSamplingY);
 
+            Av1BlockSize chromaBlockSize = blockSize.GetSubsampled(
+                colorConfig.SubSamplingX,
+                colorConfig.SubSamplingY);
+
             Span<int> blueCoefficients = this.coefficientBuffer.GetPlaneSpan(this.superblock.Index, Av1Plane.U);
             Span<int> redCoefficients = this.coefficientBuffer.GetPlaneSpan(this.superblock.Index, Av1Plane.V);
             Span<Av1EncoderTransformBlockState> blueTransformBlocks =
@@ -662,8 +666,8 @@ internal static partial class Av1IntraSuperblockEncoder
             int chromaTransformIndex = this.codedAreaChroma /
                 Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
 
-            ref Av1EncoderTransformBlockState blueState = ref blueTransformBlocks[chromaTransformIndex];
+            Span<Av1EncoderTransformBlockState> retainedBlueStates = blueTransformBlocks[chromaTransformIndex..];
-            ref Av1EncoderTransformBlockState redState = ref redTransformBlocks[chromaTransformIndex];
+            Span<Av1EncoderTransformBlockState> retainedRedStates = redTransformBlocks[chromaTransformIndex..];
             long chromaCost = 0;
             if (block.HasChroma)
             {
@@ -679,8 +683,8 @@ internal static partial class Av1IntraSuperblockEncoder
                     chromaTransformSize,
                     blueCoefficients[this.codedAreaChroma..],
                     redCoefficients[this.codedAreaChroma..],
-                    ref blueState,
+                    retainedBlueStates,
-                    ref redState,
+                    retainedRedStates,
                     ref paletteInfo,
                     out int chromaAngleDelta,
                     out byte chromaFromLumaIndex,
@@ -694,16 +698,26 @@ internal static partial class Av1IntraSuperblockEncoder
 
             // Skip suppresses coefficient syntax for the entire coding block, not one plane independently.
             // Preserve normal coefficient coding when any selected luma or chroma transform is nonempty.
-            bool allTransformsEmpty = lumaTransformEmpty &&
+            int chromaTransformSampleCount = chromaTransformSize.GetSize2d();
-                (!block.HasChroma || (blueState.EndOfBlock == 0 && redState.EndOfBlock == 0));
+            int chromaTransformBlockCount =
+                (chromaBlockSize.GetWidth() * chromaBlockSize.GetHeight()) / chromaTransformSampleCount;
+
+            int chromaStateStride =
+                chromaTransformSampleCount / Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
+
+            bool chromaTransformsEmpty = true;
+            for (int transformIndex = 0; transformIndex < chromaTransformBlockCount; transformIndex++)
+            {
+                int stateIndex = transformIndex * chromaStateStride;
+                chromaTransformsEmpty &= retainedBlueStates[stateIndex].EndOfBlock == 0 &&
+                    retainedRedStates[stateIndex].EndOfBlock == 0;
+            }
+
+            bool allTransformsEmpty = lumaTransformEmpty && (!block.HasChroma || chromaTransformsEmpty);
 
             int regularEmptyTransformRate = 0;
             if (allTransformsEmpty)
             {
-                Av1BlockSize chromaBlockSize = blockSize.GetSubsampled(
-                    colorConfig.SubSamplingX,
-                    colorConfig.SubSamplingY);
-
                 regularEmptyTransformRate = this.GetEmptyTransformRate(
                     writer,
                     this.picture.LuminanceDcSignLevelCoefficientNeighbors[tileIndex],
@@ -769,10 +783,6 @@ internal static partial class Av1IntraSuperblockEncoder
             this.codedAreaLuma += blockSize.GetWidth() * blockSize.GetHeight();
             if (block.HasChroma)
             {
-                Av1BlockSize chromaBlockSize = blockSize.GetSubsampled(
-                    colorConfig.SubSamplingX,
-                    colorConfig.SubSamplingY);
-
                 this.codedAreaChroma += chromaBlockSize.GetWidth() * chromaBlockSize.GetHeight();
             }
         }
@@ -958,8 +968,10 @@ internal static partial class Av1IntraSuperblockEncoder
             int transformWidth = transformSize.GetWidth();
             int transformHeight = transformSize.GetHeight();
             int transformSampleCount = transformSize.GetSize2d();
-            int transformIndex = 0;
+            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
             // preserves the distinct top and left contexts consumed by the next coding block in a dry run.
@@ -967,7 +979,7 @@ internal static partial class Av1IntraSuperblockEncoder
             {
                 for (int column = 0; column < blockWidth; column += transformWidth)
                 {
-                    Av1EncoderTransformBlockState state = states[transformIndex++];
+                    Av1EncoderTransformBlockState state = states[transformStateOffset];
                     byte context = Av1SymbolContextHelper.GetCoefficientContext(
                         coefficients[coefficientOffset..],
                         transformSize,
@@ -981,6 +993,7 @@ internal static partial class Av1IntraSuperblockEncoder
                         EdgeMask);
 
                     coefficientOffset += transformSampleCount;
+                    transformStateOffset += transformStateStride;
                 }
             }
         }
@@ -1685,7 +1698,9 @@ internal static partial class Av1IntraSuperblockEncoder
             }
 
             if (this.effort >= 4 &&
-                this.picture.Sequence.SequenceHeader.EnableFilterIntra)
+                this.picture.Sequence.SequenceHeader.EnableFilterIntra &&
+                blockWidth <= 32 &&
+                blockHeight <= 32)
             {
                 // 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.
@@ -2004,12 +2019,18 @@ internal static partial class Av1IntraSuperblockEncoder
                     {
                         Span<TSample> aboveStorage = workspace.GetReferenceSamples(0);
                         Span<TSample> leftStorage = workspace.GetReferenceSamples(1);
-                        this.PrepareSplitLumaReferenceSamples(
+                        this.PrepareTransformReferenceSamples(
                             reconstructionPlane,
                             blockOrigin,
+                            blockOrigin,
+                            BlockSize,
                             macroBlock,
                             transformRow,
                             transformColumn,
+                            BlockWidth,
+                            TransformSize,
+                            0,
+                            0,
                             candidateReconstruction,
                             aboveStorage,
                             leftStorage,
@@ -2167,81 +2188,89 @@ internal static partial class Av1IntraSuperblockEncoder
             return Av1RateDistortion.GetCost(this.rateMultiplier, rate, distortion);
         }
 
-        private void PrepareSplitLumaReferenceSamples(
+        private void PrepareTransformReferenceSamples(
             Buffer2DRegion<TSample> reconstructionPlane,
-            Point blockOrigin,
+            Point lumaBlockOrigin,
+            Point planeBlockOrigin,
+            Av1BlockSize blockSize,
             Av1MacroBlockD macroBlock,
             int transformRow,
             int transformColumn,
+            int planeBlockWidth,
+            Av1TransformSize transformSize,
+            int subsamplingX,
+            int subsamplingY,
             ReadOnlySpan<TSample> candidateReconstruction,
             Span<TSample> aboveStorage,
             Span<TSample> leftStorage,
             out bool hasLeft,
             out bool hasAbove)
         {
-            const Av1BlockSize BlockSize = Av1BlockSize.Block8x8;
+            int transformWidth = transformSize.GetWidth();
-            const Av1TransformSize TransformSize = Av1TransformSize.Size4x4;
+            int transformHeight = transformSize.GetHeight();
-            const int BlockWidth = 8;
+            int rowOffset = transformRow * transformHeight;
-            const int TransformWidth = 4;
+            int columnOffset = transformColumn * transformWidth;
-            int rowOffset = transformRow * TransformWidth;
+            int modeInfoRow = lumaBlockOrigin.Y >> Av1Constants.ModeInfoSizeLog2;
-            int columnOffset = transformColumn * TransformWidth;
+            int modeInfoColumn = lumaBlockOrigin.X >> Av1Constants.ModeInfoSizeLog2;
-            int modeInfoRow = blockOrigin.Y >> Av1Constants.ModeInfoSizeLog2;
-            int modeInfoColumn = blockOrigin.X >> Av1Constants.ModeInfoSizeLog2;
 
             // Internal top and left edges come from the candidate mosaic built in raster order. Edges outside
-            // the 8x8 candidate continue to read committed reconstruction, keeping unsuccessful trials isolated.
+            // the candidate continue to read committed reconstruction, keeping unsuccessful trials isolated.
             hasAbove = transformRow > 0 || macroBlock.IsUpAvailable;
             hasLeft = transformColumn > 0 || macroBlock.IsLeftAvailable;
+            int transformRow4x4 = rowOffset >> Av1Constants.ModeInfoSizeLog2;
+            int transformColumn4x4 = columnOffset >> Av1Constants.ModeInfoSizeLog2;
             bool rightAvailable =
-                modeInfoColumn + transformColumn + TransformSize.Get4x4WideCount() <
+                modeInfoColumn +
+                    ((transformColumn4x4 + transformSize.Get4x4WideCount()) << subsamplingX) <
                 macroBlock.Tile.ModeInfoColumnEnd;
 
             bool bottomAvailable =
-                modeInfoRow + transformRow + TransformSize.Get4x4HighCount() <
+                modeInfoRow +
+                    ((transformRow4x4 + transformSize.Get4x4HighCount()) << subsamplingY) <
                 macroBlock.Tile.ModeInfoRowEnd;
 
             bool hasTopRight = Av1IntraReferenceAvailability.HasTopRight(
                 this.picture.Sequence.SequenceHeader.SuperblockSize,
-                BlockSize,
+                blockSize,
                 modeInfoRow,
                 modeInfoColumn,
                 hasAbove,
                 rightAvailable,
                 Av1PartitionType.None,
-                TransformSize,
+                transformSize,
-                transformRow,
+                transformRow4x4,
-                transformColumn,
+                transformColumn4x4,
-                0,
+                subsamplingX,
-                0);
+                subsamplingY);
 
             bool hasBottomLeft = Av1IntraReferenceAvailability.HasBottomLeft(
                 this.picture.Sequence.SequenceHeader.SuperblockSize,
-                BlockSize,
+                blockSize,
                 modeInfoRow,
                 modeInfoColumn,
                 bottomAvailable,
                 hasLeft,
                 Av1PartitionType.None,
-                TransformSize,
+                transformSize,
-                transformRow,
+                transformRow4x4,
-                transformColumn,
+                transformColumn4x4,
-                0,
+                subsamplingX,
-                0);
+                subsamplingY);
 
-            Span<TSample> above = aboveStorage.Slice(1, TransformWidth * 2);
+            Span<TSample> above = aboveStorage.Slice(1, transformWidth * 2);
-            Span<TSample> left = leftStorage.Slice(1, TransformWidth * 2);
+            Span<TSample> left = leftStorage.Slice(1, transformHeight * 2);
             if (hasAbove)
             {
                 if (transformRow > 0)
                 {
                     candidateReconstruction
-                        .Slice(((rowOffset - 1) * BlockWidth) + columnOffset, TransformWidth)
+                        .Slice(((rowOffset - 1) * planeBlockWidth) + columnOffset, transformWidth)
                         .CopyTo(above);
                 }
                 else
                 {
-                    reconstructionPlane.DangerousGetRowSpan(blockOrigin.Y - 1)
+                    reconstructionPlane.DangerousGetRowSpan(planeBlockOrigin.Y - 1)
-                        .Slice(blockOrigin.X + columnOffset, TransformWidth)
+                        .Slice(planeBlockOrigin.X + columnOffset, transformWidth)
                         .CopyTo(above);
                 }
             }
@@ -2250,17 +2279,18 @@ internal static partial class Av1IntraSuperblockEncoder
             {
                 if (transformColumn > 0)
                 {
-                    for (int row = 0; row < TransformWidth; row++)
+                    for (int row = 0; row < transformHeight; row++)
                     {
-                        left[row] = candidateReconstruction[((rowOffset + row) * BlockWidth) + columnOffset - 1];
+                        left[row] = candidateReconstruction[
+                            ((rowOffset + row) * planeBlockWidth) + columnOffset - 1];
                     }
                 }
                 else
                 {
-                    for (int row = 0; row < TransformWidth; row++)
+                    for (int row = 0; row < transformHeight; row++)
                     {
                         left[row] = reconstructionPlane
-                            .DangerousGetRowSpan(blockOrigin.Y + rowOffset + row)[blockOrigin.X - 1];
+                            .DangerousGetRowSpan(planeBlockOrigin.Y + rowOffset + row)[planeBlockOrigin.X - 1];
                     }
                 }
             }
@@ -2268,12 +2298,12 @@ internal static partial class Av1IntraSuperblockEncoder
             int midpoint = 128 << (this.bitDepth.GetBitCount() - 8);
             if (!hasAbove)
             {
-                above[..TransformWidth].Fill(hasLeft ? left[0] : TOperator.CreateSample(midpoint - 1));
+                above[..transformWidth].Fill(hasLeft ? left[0] : TOperator.CreateSample(midpoint - 1));
             }
 
             if (!hasLeft)
             {
-                left[..TransformWidth].Fill(hasAbove ? above[0] : TOperator.CreateSample(midpoint + 1));
+                left[..transformHeight].Fill(hasAbove ? above[0] : TOperator.CreateSample(midpoint + 1));
             }
 
             if (hasTopRight)
@@ -2282,42 +2312,42 @@ internal static partial class Av1IntraSuperblockEncoder
                 {
                     candidateReconstruction
                         .Slice(
-                            ((rowOffset - 1) * BlockWidth) + columnOffset + TransformWidth,
+                            ((rowOffset - 1) * planeBlockWidth) + columnOffset + transformWidth,
-                            TransformWidth)
+                            transformWidth)
-                        .CopyTo(above[TransformWidth..]);
+                        .CopyTo(above[transformWidth..]);
                 }
                 else
                 {
-                    reconstructionPlane.DangerousGetRowSpan(blockOrigin.Y - 1)
+                    reconstructionPlane.DangerousGetRowSpan(planeBlockOrigin.Y - 1)
-                        .Slice(blockOrigin.X + columnOffset + TransformWidth, TransformWidth)
+                        .Slice(planeBlockOrigin.X + columnOffset + transformWidth, transformWidth)
-                        .CopyTo(above[TransformWidth..]);
+                        .CopyTo(above[transformWidth..]);
                 }
             }
             else
             {
-                above[TransformWidth..].Fill(above[TransformWidth - 1]);
+                above[transformWidth..].Fill(above[transformWidth - 1]);
             }
 
             if (hasBottomLeft)
             {
-                for (int row = TransformWidth; row < TransformWidth * 2; row++)
+                for (int row = transformHeight; row < transformHeight * 2; row++)
                 {
                     left[row] = reconstructionPlane
-                        .DangerousGetRowSpan(blockOrigin.Y + rowOffset + row)[blockOrigin.X - 1];
+                        .DangerousGetRowSpan(planeBlockOrigin.Y + rowOffset + row)[planeBlockOrigin.X - 1];
                 }
             }
             else
             {
-                left[TransformWidth..].Fill(left[TransformWidth - 1]);
+                left[transformHeight..].Fill(left[transformHeight - 1]);
             }
 
             // Only an interior transform corner belongs to decision scratch. Boundary corners continue
             // to read the already reconstructed neighboring block so candidate trials remain isolated.
             TSample corner = hasAbove && hasLeft
                 ? transformRow > 0 && transformColumn > 0
-                    ? candidateReconstruction[((rowOffset - 1) * BlockWidth) + columnOffset - 1]
+                    ? candidateReconstruction[((rowOffset - 1) * planeBlockWidth) + columnOffset - 1]
-                    : reconstructionPlane.DangerousGetRowSpan(blockOrigin.Y + rowOffset - 1)[
+                    : reconstructionPlane.DangerousGetRowSpan(planeBlockOrigin.Y + rowOffset - 1)[
-                        blockOrigin.X + columnOffset - 1]
+                        planeBlockOrigin.X + columnOffset - 1]
                 : hasAbove
                     ? above[0]
                     : hasLeft

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

@@ -1914,84 +1914,78 @@ internal partial class Av1TileWriter
         int maximumUnitBlocksWide = Math.Min(maximumUnitBlockSize.Get4x4WideCount(), maximumBlocksWide);
         int maximumUnitBlocksHigh = Math.Min(maximumUnitBlockSize.Get4x4HighCount(), maximumBlocksHigh);
 
-        // Chroma follows the same bounded-region order after scaling both the block and frame edges to its plane.
+        int codedAreaStart = entropyCodingContext.CodedAreaSuperblockUv;
-        for (int regionRow = 0; regionRow < maximumBlocksHigh; regionRow += maximumUnitBlocksHigh)
+        int codedAreaEnd = codedAreaStart;
+
+        // AV1 completes every transform in one chroma plane before advancing to the other plane.
+        // Both planes use the same coded-area positions because their transform geometry is identical.
+        for (int planeIndex = 0; planeIndex < 2; planeIndex++)
         {
-            int unitHeight = Math.Min(maximumUnitBlocksHigh + regionRow, maximumBlocksHigh);
+            bool isBluePlane = planeIndex == 0;
-            for (int regionColumn = 0; regionColumn < maximumBlocksWide; regionColumn += maximumUnitBlocksWide)
+            Span<int> planeCoefficients = isBluePlane ? blueCoefficients : redCoefficients;
+            Span<Av1EncoderTransformBlockState> planeTransformBlocks =
+                isBluePlane ? blueTransformBlocks : redTransformBlocks;
+
+            Av1NeighborArrayUnit<byte> coefficientNeighbors =
+                isBluePlane ? cb_dc_sign_level_coeff_na : cr_dc_sign_level_coeff_na;
+
+            int codedArea = codedAreaStart;
+            for (int regionRow = 0; regionRow < maximumBlocksHigh; regionRow += maximumUnitBlocksHigh)
             {
-                int unitWidth = Math.Min(maximumUnitBlocksWide + regionColumn, maximumBlocksWide);
+                int unitHeight = Math.Min(maximumUnitBlocksHigh + regionRow, maximumBlocksHigh);
-                for (int blockRow = regionRow; blockRow < unitHeight; blockRow += transformBlockHeight)
+                for (int regionColumn = 0; regionColumn < maximumBlocksWide; regionColumn += maximumUnitBlocksWide)
                 {
-                    for (int blockColumn = regionColumn; blockColumn < unitWidth; blockColumn += transformBlockWidth)
+                    int unitWidth = Math.Min(maximumUnitBlocksWide + regionColumn, maximumBlocksWide);
+                    for (int blockRow = regionRow; blockRow < unitHeight; blockRow += transformBlockHeight)
                     {
-                        int transformStateIndex = entropyCodingContext.CodedAreaSuperblockUv /
+                        for (int blockColumn = regionColumn; blockColumn < unitWidth; blockColumn += transformBlockWidth)
-                            Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
+                        {
-                        ref Av1EncoderTransformBlockState blueTransformBlock = ref blueTransformBlocks[transformStateIndex];
+                            int transformStateIndex =
-                        ref Av1EncoderTransformBlockState redTransformBlock = ref redTransformBlocks[transformStateIndex];
+                                codedArea / Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount;
-                        Point chromaOrigin = chromaBlockOrigin + new Size(
+
-                            blockColumn << Av1Constants.ModeInfoSizeLog2,
+                            ref Av1EncoderTransformBlockState transformBlock =
-                            blockRow << Av1Constants.ModeInfoSizeLog2);
+                                ref planeTransformBlocks[transformStateIndex];
-
+
-                        // U and V share transform geometry and type while retaining independent coefficient and EOB state.
+                            Point chromaOrigin = chromaBlockOrigin + new Size(
-                        Span<int> coefficients = blueCoefficients[entropyCodingContext.CodedAreaSuperblockUv..];
+                                blockColumn << Av1Constants.ModeInfoSizeLog2,
-                        Av1TransformBlockContext blockContext = GetTransformBlockContexts(
+                                blockRow << Av1Constants.ModeInfoSizeLog2);
-                            Av1ComponentType.Chroma,
+
-                            cb_dc_sign_level_coeff_na,
+                            Span<int> coefficients = planeCoefficients[codedArea..];
-                            chromaOrigin,
+                            Av1TransformBlockContext blockContext = GetTransformBlockContexts(
-                            chromaBlockSize,
+                                Av1ComponentType.Chroma,
-                            chromaTransformSize);
+                                coefficientNeighbors,
-
+                                chromaOrigin,
-                        Av1TransformType chromaTransformType = blueTransformBlock.TransformType;
+                                chromaBlockSize,
-                        int culLevelCb = writer.WriteCoefficients(
+                                chromaTransformSize);
-                            chromaTransformSize,
+
-                            chromaTransformType,
+                            int culLevel = writer.WriteCoefficients(
-                            intraLumaDir,
+                                chromaTransformSize,
-                            coefficients,
+                                transformBlock.TransformType,
-                            Av1ComponentType.Chroma,
+                                intraLumaDir,
-                            blockContext,
+                                coefficients,
-                            blueTransformBlock.EndOfBlock,
+                                Av1ComponentType.Chroma,
-                            frameHeader.UseReducedTransformSet,
+                                blockContext,
-                            blk_ptr.FilterIntraMode,
+                                transformBlock.EndOfBlock,
-                            usesInterTransformSet);
+                                frameHeader.UseReducedTransformSet,
-
+                                blk_ptr.FilterIntraMode,
-                        coefficients = redCoefficients[entropyCodingContext.CodedAreaSuperblockUv..];
+                                usesInterTransformSet);
-                        blockContext = GetTransformBlockContexts(
+
-                            Av1ComponentType.Chroma,
+                            coefficientNeighbors.UnitModeWrite(
-                            cr_dc_sign_level_coeff_na,
+                                (byte)culLevel,
-                            chromaOrigin,
+                                chromaOrigin,
-                            chromaBlockSize,
+                                new Size(transformWidth, transformHeight),
-                            chromaTransformSize);
+                                Av1NeighborArrayUnit<byte>.UnitMask.Top | Av1NeighborArrayUnit<byte>.UnitMask.Left);
-
+
-                        int culLevelCr = writer.WriteCoefficients(
+                            codedArea += transformWidth * transformHeight;
-                            chromaTransformSize,
+                        }
-                            chromaTransformType,
-                            intraLumaDir,
-                            coefficients,
-                            Av1ComponentType.Chroma,
-                            blockContext,
-                            redTransformBlock.EndOfBlock,
-                            frameHeader.UseReducedTransformSet,
-                            blk_ptr.FilterIntraMode,
-                            usesInterTransformSet);
-
-                        cb_dc_sign_level_coeff_na.UnitModeWrite(
-                            (byte)culLevelCb,
-                            chromaOrigin,
-                            new Size(transformWidth, transformHeight),
-                            Av1NeighborArrayUnit<byte>.UnitMask.Top | Av1NeighborArrayUnit<byte>.UnitMask.Left);
-
-                        cr_dc_sign_level_coeff_na.UnitModeWrite(
-                            (byte)culLevelCr,
-                            chromaOrigin,
-                            new Size(transformWidth, transformHeight),
-                            Av1NeighborArrayUnit<byte>.UnitMask.Top | Av1NeighborArrayUnit<byte>.UnitMask.Left);
-
-                        entropyCodingContext.CodedAreaSuperblockUv += transformWidth * transformHeight;
                     }
                 }
             }
+
+            codedAreaEnd = codedArea;
         }
+
+        entropyCodingContext.CodedAreaSuperblockUv = codedAreaEnd;
     }
 
     /// <summary>

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

@@ -277,7 +277,7 @@ public class Av1EncoderFrameTests
     [Fact]
     public void EncodeEffortTenSelectsThirtyTwoByThirtyTwoBlocks()
     {
-        const int Size = 64;
+        const int Size = 32;
         using Image<Rgba32> source = new(Size, Size);
         for (int y = 0; y < Size; y++)
         {
@@ -297,6 +297,49 @@ public class Av1EncoderFrameTests
             qIndex: 4,
             effort: 10);
 
+        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 < 8; modeInfoY++)
+        {
+            for (int modeInfoX = 0; modeInfoX < 8; modeInfoX++)
+            {
+                Assert.Equal(
+                    Av1BlockSize.Block32x32,
+                    frameInfo.GetModeInfoAt(new Point(modeInfoX, modeInfoY)).BlockSize);
+            }
+        }
+
+        Assert.Equal(new Size(Size, Size), decoded.Size);
+    }
+
+    [Theory]
+    [InlineData(Yuv400)]
+    [InlineData(Yuv444)]
+    public void EncodeEffortTenSelectsSixtyFourBySixtyFourBlock(int colorFormatValue)
+    {
+        const int Size = 64;
+        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();
+        _ = Av1FrameEncoder.Encode(
+            Configuration.Default,
+            source.Frames.RootFrame,
+            stream,
+            CreateColorConfig(Av1BitDepth.EightBit, colorFormat),
+            qIndex: 4,
+            effort: 10);
+
         byte[] payload = stream.ToArray();
         using Av1Decoder decoder = new(Configuration.Default);
         using Image<Rgba32> decoded = decoder.Decode<Rgba32>(payload);
@@ -306,7 +349,7 @@ public class Av1EncoderFrameTests
             for (int modeInfoX = 0; modeInfoX < 16; modeInfoX++)
             {
                 Assert.Equal(
-                    Av1BlockSize.Block32x32,
+                    Av1BlockSize.Block64x64,
                     frameInfo.GetModeInfoAt(new Point(modeInfoX, modeInfoY)).BlockSize);
             }
         }

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

@@ -623,10 +623,9 @@ public class Av1TransformBlockEncoderTests
 
             Assert.Equal(Av1EncoderModeDecisionWorkspace<ushort>.MaximumSampleCount, modeWorkspace.GetCandidateReconstruction(1).Length);
             Assert.Equal(Av1EncoderModeDecisionWorkspace<ushort>.MaximumSampleCount, modeWorkspace.GetCandidateCoefficients(1).Length);
-            Assert.Equal(
+
-                Av1ChromaFromLumaContext.BufferLine *
+            // CfL is unavailable above 32x32, so its scratch remains fixed while larger partitions are enabled.
-                    Av1EncoderModeDecisionWorkspace<ushort>.MaximumBlockDimension,
+            Assert.Equal(Av1ChromaFromLumaContext.BufferLength, modeWorkspace.ChromaFromLumaSamples.Length);
-                modeWorkspace.ChromaFromLumaSamples.Length);
 
             Assert.Equal(Av1ChromaFromLumaMath.AlphaCandidateCount, modeWorkspace.GetChromaFromLumaRates(1).Length);
             Assert.Equal(Av1ChromaFromLumaMath.AlphaCandidateCount, modeWorkspace.GetChromaFromLumaDistortions(1).Length);