From 41691eeb80539eb886a5734500e524c2c7e62149 Mon Sep 17 00:00:00 2001 From: James Jackson-South Date: Wed, 2 Sep 2026 18:25:34 +1000 Subject: [PATCH] Encode fixed AV1 intra superblocks --- HEIF_IMPLEMENTATION_PLAN.md | 6 +- .../Av1IntraSuperblockEncoder.Operator.cs | 161 +++++++ .../Av1/Pipeline/Av1IntraSuperblockEncoder.cs | 276 ++++++++++++ .../Av1/Av1IntraSuperblockEncoderTests.cs | 401 ++++++++++++++++++ 4 files changed, 841 insertions(+), 3 deletions(-) create mode 100644 src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.Operator.cs create mode 100644 src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.cs create mode 100644 tests/ImageSharp.Tests/Formats/Heif/Av1/Av1IntraSuperblockEncoderTests.cs diff --git a/HEIF_IMPLEMENTATION_PLAN.md b/HEIF_IMPLEMENTATION_PLAN.md index b382ee6cec..167485b247 100644 --- a/HEIF_IMPLEMENTATION_PLAN.md +++ b/HEIF_IMPLEMENTATION_PLAN.md @@ -819,11 +819,11 @@ Encoder verification contract: ### 6. Build the complete AV1 frame encoder - [~] SIMD-first RGB-to-native-plane conversion now feeds eight-bit and high-bit-depth bordered AV1 source frames directly, preserving ImageSharp's arbitrary packed-pixel input contract without an intermediate full-frame native-plane copy. -- [~] Forward transform families, transform workspace, and an allocation-free DC intra block boundary exist locally. For eight-bit and high-bit-depth samples, the composed boundary now follows current libaom's encoder order: predict into the reconstruction plane, subtract prediction from source, transform, quantize into separate qcoeff and dqcoeff storage, retain EOB and transform type, and inverse-transform only when EOB is nonzero so later blocks consume decoder-identical references. Prediction and subtraction retain their SIMD-first operators, independent source and reconstruction strides are preserved, and no frame-sized or per-block buffer is introduced. The block boundary consumes the real bordered encoder-plane regions and indexes their one-segment owner directly; this preserves physical row strides without a row copy and avoids the per-call enumerator allocation exposed by the initial array-only test. One reusable 61 KiB allocator owner supplies tightly packed residual, aligned transform-coefficient, dequantized-coefficient, and transform scratch spans across transform blocks; quantized coefficients write directly to the retained frame coefficient owner instead of being duplicated. Stage-by-stage scalar-oracle, physical-border, retained-syntax, and zero-allocation coverage passes 6 of 6 through direct net11 VSTest in Release. Frame traversal still needs to select blocks, gather contiguous left references, and supply coefficient-owner slices. +- [~] Forward transform families, transform workspace, and an allocation-free DC intra block boundary exist locally. For eight-bit and high-bit-depth samples, the composed boundary now follows current libaom's encoder order: predict into the reconstruction plane, subtract prediction from source, transform, quantize into separate qcoeff and dqcoeff storage, retain EOB and transform type, and inverse-transform only when EOB is nonzero so later blocks consume decoder-identical references. Prediction and subtraction retain their SIMD-first operators, independent source and reconstruction strides are preserved, and no frame-sized or per-block buffer is introduced. The block boundary consumes the real bordered encoder-plane regions and indexes their one-segment owner directly; this preserves physical row strides without a row copy and avoids the per-call enumerator allocation exposed by the initial array-only test. One reusable 61 KiB allocator owner supplies tightly packed residual, aligned transform-coefficient, dequantized-coefficient, and transform scratch spans across transform blocks; quantized coefficients write directly to the retained frame coefficient owner instead of being duplicated. A fixed 8x8 DC-intra superblock baseline now traverses the same recursive preorder and frame-edge pruning as the tile writer, gathers left references into that reusable block workspace, writes luma and chroma coefficient-owner slices in the writer's exact consumption order, and updates the caller-owned reconstruction planes for subsequent predictions. Stage-by-stage scalar-oracle, physical-border, retained-syntax, superblock-to-writer synchronization, high-bit-depth precision, and steady-state zero-allocation coverage passes 8 of 8 through direct net11 VSTest in Release. This is a legal fixed baseline, not complete partition or mode analysis. - [~] Symbol writer, coefficient writer, and tile writer fragments exist locally. - [~] A non-owning encoder-frame view now separates visible conversion regions from coded regions and performs complete left, top, right, bottom, and corner extension across each bordered plane. Current libaom uses 8-sample-aligned coded dimensions, a 32-sample-aligned luma stride with chroma stride derived from it, and a 64-pixel luma border for non-resized all-intra encoding. One operation-ready frame owner now rents the aligned Y, U, and V storage contiguously, exposes non-owning `Buffer2D` plane views, and returns the rent exactly once. A 4K 4:2:0 frame occupies about 13.0 MiB at 8-bit or 26.0 MiB at 10/12-bit; source and reconstruction therefore remain distinct frame owners rather than adding a full-frame copy. The corrected tests use this real ownership path and verify the exact 54 KiB 64x64 4:2:0 rent. The frame-encoder boundary converts packed pixels directly into the source owner before extension; the containing encode operation still needs to instantiate matching source and reconstruction owners with ordinary `using` lifetimes. - [~] Temporal delimiter, sequence header, frame header, and combined-frame tile-group writing exist locally. The remaining required metadata, padding, and encoder-wide syntax paths are not complete. -- [ ] Implement superblock and partition analysis for every permitted block size and partition. +- [~] Implement superblock and partition analysis for every permitted block size and partition. The current baseline deliberately splits every in-frame node to 8x8 blocks and records decisions in current-libaom writer preorder; block-size selection and non-split partition analysis remain. - [ ] Implement intra mode search, chroma mode search, palette, filter intra, chroma-from-luma, and intra-block copy decisions. - [ ] Implement inter mode search for bounded sequences, including reference selection and the decoder-supported inter tools. - [~] Current-libaom `av1_quantize_fp_no_qmatrix` arithmetic is implemented as a closed generic forward-quantizer family with Vector512, Vector256, Vector128, and scalar paths, raster-order output, coded 64-point coefficient limits, and scan-order EOB selection. Transform search, coefficient optimization, and lossless behavior remain. @@ -831,7 +831,7 @@ Encoder verification contract: - [~] The tile writer now publishes one packed coefficient context per covered 4x4 edge unit and derives luma/chroma skip plus DC-sign contexts from the complete transform edges using current-libaom units. Partition, transform, and coefficient neighbor state now retains only the above and left context regions used by current libaom; the unused third top-left region, its granularity state, and its unused sentinel are removed. One clean allocation contains the two active edges, and the exact requested length plus exactly-once return pass with the complete 77-case coefficient and entropy class in direct net11 VSTest Release. Complete tile traversal, initialized picture state, and verified CDF update behavior remain. - [~] Encoder mode information now uses a frame-owned integer alias grid over a packed 8-byte value allocation, matching current libaom's `mi_grid_base` and `mi_alloc` relationship without a managed object or reference per 4x4 entry. The visible dimensions are aligned to eight luma samples, the grid stride and allocated row count are aligned to 32 mode-information units, and optional 8x8 allocation granularity reduces the value store in both dimensions exactly as current libaom does. One clean ImageSharp byte owner contains both independently typed regions, reducing libaom's two allocation lifetimes to one without a copy. At 4K, the 4x4 layout occupies about 6.0 MiB in total; the 8x8 layout occupies about 3.0 MiB. Exact geometry, clean allocation, typed lengths, aligned mapping, untouched row padding, and exactly-once return pass 4 of 4 direct net11 VSTest cases in Release. Every coded 4x4 cell covered by square, rectangular, or clipped edge blocks maps to its owning allocation entry before context-dependent symbols are written. Packed syntax, relative neighbor lookup, full block mapping, writer traversal, entropy, and OBU coverage pass 1,947 of 1,947 direct net11 VSTest cases in Release; complete mode decision still remains. - [~] The final-block decision workspace uses one reusable 10.3 KiB ImageSharp allocator owner. It contains 1,024 explicitly packed 10-byte final-block entries and the 341 preorder partition bytes required by a complete 128x128-through-8x8 quadtree, replacing separate managed arrays. Construction and the explicit per-superblock reset initialize every syntax field, including the nonzero sentinel that disables filter-intra prediction; pooled palette, quantizer, prediction, and partition bytes cannot leak into the next decision pass. Exact allocation, size, initialization, reset, return, repeated-run, writer, entropy, and OBU coverage pass 1,957 of 1,957 direct net11 VSTest cases in Release; complete mode decision still remains. -- [~] Finalized transform coefficients and packed EOB/type state now use raster-ordered, per-superblock plane segments matching current libaom's coefficient-pool geometry. One ImageSharp allocator owner replaces libaom's separate coefficient, EOB, and entropy-context allocations while preserving the full 1024 luma and 256-per-chroma 4x4 state capacity of a 128x128 4:2:0 superblock. The composed intra block boundary can populate the owner's quantized coefficient and state slices while updating the caller-owned reconstruction plane directly; mode-decision traversal still needs to select and invoke it. +- [~] Finalized transform coefficients and packed EOB/type state now use raster-ordered, per-superblock plane segments matching current libaom's coefficient-pool geometry. One ImageSharp allocator owner replaces libaom's separate coefficient, EOB, and entropy-context allocations while preserving the full 1024 luma and 256-per-chroma 4x4 state capacity of a 128x128 4:2:0 superblock. The fixed 8x8 DC-intra traversal populates the owner's quantized coefficient and state slices while updating the caller-owned reconstruction plane directly, and a real tile-writer integration check proves that both sides consume identical luma and chroma areas. Complete mode decision still remains. - [~] Tile partition writing now follows current libaom's recursive `write_modes_sb` preorder traversal and `update_ext_partition_context` edge updates directly. Bottom-edge blocks use the horizontal-alike partition CDF and right-edge blocks use the vertical-alike CDF; byte-exact regressions cover both paths after the previous calls were found reversed. Lossless chroma-from-luma availability now uses the subsampled plane block size shared with the decoder instead of the lossy 32x32 limit, preserving the correct UV-mode alphabet for each segment. The obsolete SVT-derived global geometry catalog and its unimplemented lookup are removed; transform geometry is derived in libaom's bounded 64x64 residual order, fixed intra transform-size symbols use the reference depth and neighbor contexts, and each derived transform size is persisted to the frame-owned mode information before the entropy snapshot and coefficient traversal consume it. Frame-edge and segmentation syntax use mode-information units, and 128x128 CDEF units use libaom's 0-to-3 indexing and first-block strength ownership. The focused transform-state regression passes 3 of 3 direct net11 VSTest cases in Release. Writer, entropy, and OBU coverage passes 1,957 of 1,957 direct net11 VSTest cases in Release, with 20 of 20 focused encoder and decoder chroma-from-luma cases. Partition and mode analysis still need to populate these retained decisions; variable inter-transform syntax remains part of later inter-frame support. - [ ] Implement legal deblocking, CDEF, restoration, super-resolution, and film-grain signaling decisions. - [~] The coefficient symbol encoder now reuses tile-lifetime level and context workspaces instead of allocating per transform, defers both coefficient rents until the first nonzero transform block, and disposes all tile scratch independently from the detached encoded bytes. Its range coder matches current libaom's 64-bit coding window, bulk big-endian byte flush, and backward carry propagation while using one byte of allocator scratch per estimated output byte instead of the former 16-bit pre-carry storage. The reference-type symbol encoder is passed normally through tile traversal, and the operation boundary owns the allocator-backed item payload stream for exactly one synchronous encode. Every remaining encoder fragment must be audited before it becomes active. diff --git a/src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.Operator.cs b/src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.Operator.cs new file mode 100644 index 0000000000..2af6b814af --- /dev/null +++ b/src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.Operator.cs @@ -0,0 +1,161 @@ +// Copyright (c) Six Labors. +// Licensed under the Six Labors Split License. + +using System.Runtime.InteropServices; +using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling; +using SixLabors.ImageSharp.Formats.Heif.Av1.Transform; +using SixLabors.ImageSharp.Memory; + +namespace SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline; + +/// +/// Defines the sample-storage operations used by fixed intra superblock traversal. +/// +internal static partial class Av1IntraSuperblockEncoder +{ + /// + /// Defines type-specific block encoding without coupling traversal to sample storage width. + /// + /// The native unsigned sample storage type. + private interface IBlockEncodingOperator + where TSample : unmanaged + { + /// + /// Gets temporary contiguous storage for left reference samples. + /// + /// The reusable residual workspace. + /// The number of reference samples. + /// The writable reference span. + public static abstract Span GetLeftReference(Span residual, int length); + + /// + /// Encodes and reconstructs one DC intra transform block. + /// + /// The reusable block workspace. + /// The coded source plane. + /// The coded reconstruction plane. + /// The transform-block origin in plane samples. + /// The top reference samples. + /// The left reference samples. + /// Whether the left reference is available. + /// Whether the top reference is available. + /// The retained entropy-coding coefficients. + /// The transform dimensions. + /// The effective segment quantizer index. + /// The plane DC quantizer adjustment. + /// The plane AC quantizer adjustment. + /// The component plane containing the block. + /// The coded sample bit depth. + /// The retained transform state. + public static abstract void Encode( + Av1EncoderBlockWorkspace workspace, + Buffer2DRegion source, + Buffer2DRegion reconstruction, + Point blockOrigin, + ReadOnlySpan above, + ReadOnlySpan left, + bool hasLeft, + bool hasAbove, + Span quantizedCoefficients, + Av1TransformSize transformSize, + int qIndex, + int dcDeltaQ, + int acDeltaQ, + Av1Plane plane, + Av1BitDepth bitDepth, + ref Av1EncoderTransformBlockState state); + } + + /// + /// Encodes blocks stored as eight-bit samples. + /// + private readonly struct ByteOperator : IBlockEncodingOperator + { + /// + public static Span GetLeftReference(Span residual, int length) + => MemoryMarshal.AsBytes(residual)[..length]; + + /// + public static void Encode( + Av1EncoderBlockWorkspace workspace, + Buffer2DRegion source, + Buffer2DRegion reconstruction, + Point blockOrigin, + ReadOnlySpan above, + ReadOnlySpan left, + bool hasLeft, + bool hasAbove, + Span quantizedCoefficients, + Av1TransformSize transformSize, + int qIndex, + int dcDeltaQ, + int acDeltaQ, + Av1Plane plane, + Av1BitDepth bitDepth, + ref Av1EncoderTransformBlockState state) + => Av1TransformBlockEncoder.EncodeIntraDcLossy( + workspace, + source, + reconstruction, + blockOrigin, + above, + left, + hasLeft, + hasAbove, + quantizedCoefficients, + transformSize, + Av1TransformType.DctDct, + qIndex, + dcDeltaQ, + acDeltaQ, + plane, + ref state); + } + + /// + /// Encodes blocks stored as high-bit-depth samples. + /// + private readonly struct UInt16Operator : IBlockEncodingOperator + { + /// + public static Span GetLeftReference(Span residual, int length) + => MemoryMarshal.Cast(residual)[..length]; + + /// + public static void Encode( + Av1EncoderBlockWorkspace workspace, + Buffer2DRegion source, + Buffer2DRegion reconstruction, + Point blockOrigin, + ReadOnlySpan above, + ReadOnlySpan left, + bool hasLeft, + bool hasAbove, + Span quantizedCoefficients, + Av1TransformSize transformSize, + int qIndex, + int dcDeltaQ, + int acDeltaQ, + Av1Plane plane, + Av1BitDepth bitDepth, + ref Av1EncoderTransformBlockState state) + => Av1TransformBlockEncoder.EncodeIntraDcLossy( + workspace, + source, + reconstruction, + blockOrigin, + above, + left, + hasLeft, + hasAbove, + quantizedCoefficients, + transformSize, + Av1TransformType.DctDct, + qIndex, + dcDeltaQ, + acDeltaQ, + plane, + bitDepth, + ref state); + } +} diff --git a/src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.cs b/src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.cs new file mode 100644 index 0000000000..df9453217c --- /dev/null +++ b/src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.cs @@ -0,0 +1,276 @@ +// Copyright (c) Six Labors. +// Licensed under the Six Labors Split License. + +using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit; +using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction; +using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling; +using SixLabors.ImageSharp.Formats.Heif.Av1.Transform; +using SixLabors.ImageSharp.Memory; + +namespace SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline; + +/// +/// Builds fixed DC intra decisions and reconstructed samples for one AV1 superblock. +/// +internal static partial class Av1IntraSuperblockEncoder +{ + /// + /// Encodes one superblock stored as eight-bit samples. + /// + /// The coded source frame. + /// The reconstructed frame updated by the block transforms. + /// The frame coding and mode-information state. + /// The reusable partition and final-block decisions. + /// The frame-owned quantized coefficient and transform state. + /// The reusable block arithmetic workspace. + public static void Encode( + Av1EncoderFrame source, + Av1EncoderFrame reconstruction, + Av1PictureControlSet picture, + Av1Superblock superblock, + Av1EncoderCoefficientBuffer coefficientBuffer, + Av1EncoderBlockWorkspace blockWorkspace) + => Encode(source, reconstruction, picture, superblock, coefficientBuffer, blockWorkspace); + + /// + /// Encodes one superblock stored as high-bit-depth samples. + /// + /// The coded source frame. + /// The reconstructed frame updated by the block transforms. + /// The frame coding and mode-information state. + /// The reusable partition and final-block decisions. + /// The frame-owned quantized coefficient and transform state. + /// The reusable block arithmetic workspace. + public static void Encode( + Av1EncoderFrame source, + Av1EncoderFrame reconstruction, + Av1PictureControlSet picture, + Av1Superblock superblock, + Av1EncoderCoefficientBuffer coefficientBuffer, + Av1EncoderBlockWorkspace blockWorkspace) + => Encode(source, reconstruction, picture, superblock, coefficientBuffer, blockWorkspace); + + private static void Encode( + Av1EncoderFrame source, + Av1EncoderFrame reconstruction, + Av1PictureControlSet picture, + Av1Superblock superblock, + Av1EncoderCoefficientBuffer coefficientBuffer, + Av1EncoderBlockWorkspace blockWorkspace) + where TSample : unmanaged + where TOperator : struct, IBlockEncodingOperator + { + int superblockSize = picture.Sequence.SequenceHeader.SuperblockSize.GetWidth(); + Point superblockOrigin = new( + (superblock.Index % coefficientBuffer.SuperblockColumnCount) * superblockSize, + (superblock.Index / coefficientBuffer.SuperblockColumnCount) * superblockSize); + + superblock.Workspace.Reset(); + Traversal traversal = new( + source.CodedView, + reconstruction.CodedView, + picture, + superblock, + coefficientBuffer, + blockWorkspace); + + traversal.EncodePartitionTree(superblockOrigin, picture.Sequence.SequenceHeader.SuperblockSize); + } + + /// + /// Retains the stack-only state shared by recursive partition and final-block traversal. + /// + /// The native unsigned sample storage type. + /// The type-specific block encoding operations. + private ref struct Traversal + where TSample : unmanaged + where TOperator : struct, IBlockEncodingOperator + { + private readonly Av1EncoderFrame.PlanarView source; + private readonly Av1EncoderFrame.PlanarView reconstruction; + private readonly Av1PictureControlSet picture; + private readonly Av1Superblock superblock; + private readonly Av1EncoderBlockWorkspace blockWorkspace; + private readonly ObuQuantizationParameters quantization; + private readonly Av1BitDepth bitDepth; + private readonly Span lumaCoefficients; + private readonly Span blueCoefficients; + private readonly Span redCoefficients; + private readonly Span lumaTransformBlocks; + private readonly Span blueTransformBlocks; + private readonly Span redTransformBlocks; + private int partitionIndex; + private int finalBlockIndex; + private int codedAreaLuma; + private int codedAreaChroma; + + public Traversal( + Av1EncoderFrame.PlanarView source, + Av1EncoderFrame.PlanarView reconstruction, + Av1PictureControlSet picture, + Av1Superblock superblock, + Av1EncoderCoefficientBuffer coefficientBuffer, + Av1EncoderBlockWorkspace blockWorkspace) + { + this.source = source; + this.reconstruction = reconstruction; + this.picture = picture; + this.superblock = superblock; + this.blockWorkspace = blockWorkspace; + this.quantization = picture.Parent.FrameHeader.QuantizationParameters; + this.bitDepth = picture.Sequence.SequenceHeader.ColorConfig.BitDepth; + this.lumaCoefficients = coefficientBuffer.GetPlaneSpan(superblock.Index, Av1Plane.Y); + this.blueCoefficients = coefficientBuffer.GetPlaneSpan(superblock.Index, Av1Plane.U); + this.redCoefficients = coefficientBuffer.GetPlaneSpan(superblock.Index, Av1Plane.V); + this.lumaTransformBlocks = coefficientBuffer.GetTransformBlockSpan(superblock.Index, Av1Plane.Y); + this.blueTransformBlocks = coefficientBuffer.GetTransformBlockSpan(superblock.Index, Av1Plane.U); + this.redTransformBlocks = coefficientBuffer.GetTransformBlockSpan(superblock.Index, Av1Plane.V); + this.partitionIndex = 0; + this.finalBlockIndex = 0; + this.codedAreaLuma = 0; + this.codedAreaChroma = 0; + } + + public void EncodePartitionTree(Point blockOrigin, Av1BlockSize blockSize) + { + Av1EncoderCommon common = this.picture.Parent.Common; + Point modeInfoPosition = blockOrigin >> Av1Constants.ModeInfoSizeLog2; + if (modeInfoPosition.Y >= common.ModeInfoRowCount || modeInfoPosition.X >= common.ModeInfoColumnCount) + { + return; + } + + if (blockSize == Av1BlockSize.Block8x8) + { + this.superblock.CodingUnitPartitionTypes[this.partitionIndex++] = (byte)Av1PartitionType.None; + this.EncodeFinalBlock(blockOrigin, modeInfoPosition); + return; + } + + this.superblock.CodingUnitPartitionTypes[this.partitionIndex++] = (byte)Av1PartitionType.Split; + Av1BlockSize subSize = Av1PartitionType.Split.GetBlockSubSize(blockSize); + int halfBlockSize = blockSize.GetWidth() >> 1; + + // The preorder and out-of-frame pruning match tile writing, so one reusable decision workspace is sufficient. + this.EncodePartitionTree(blockOrigin, subSize); + this.EncodePartitionTree(blockOrigin + new Size(halfBlockSize, 0), subSize); + this.EncodePartitionTree(blockOrigin + new Size(0, halfBlockSize), subSize); + this.EncodePartitionTree(blockOrigin + new Size(halfBlockSize, halfBlockSize), subSize); + } + + private void EncodeFinalBlock(Point blockOrigin, Point modeInfoPosition) + { + const Av1BlockSize BlockSize = Av1BlockSize.Block8x8; + const Av1TransformSize LumaTransformSize = Av1TransformSize.Size8x8; + int qIndex = this.quantization.QIndex[0]; + ref Av1MacroBlockModeInfo modeInfo = ref this.picture.GetMacroBlockModeInfo(modeInfoPosition); + modeInfo.Block = new Av1EncoderBlockModeInfo + { + BlockSize = BlockSize, + PartitionType = Av1PartitionType.None, + SegmentId = 0, + TransformSize = LumaTransformSize, + Mode = Av1PredictionMode.DC, + UvMode = Av1ChromaPredictionMode.DC + }; + + modeInfo.CdefStrength = 0; + ref Av1EncoderBlockStruct block = ref this.superblock.FinalBlocks[this.finalBlockIndex++]; + block.HasChroma = !this.source.IsMonochrome; + block.QuantizationIndex = qIndex; + block.SegmentId = 0; + + int lumaTransformIndex = this.codedAreaLuma / + Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount; + + ref Av1EncoderTransformBlockState lumaState = ref this.lumaTransformBlocks[lumaTransformIndex]; + this.EncodePlaneBlock( + Av1Plane.Y, + blockOrigin, + LumaTransformSize, + this.lumaCoefficients[this.codedAreaLuma..], + ref lumaState); + + this.codedAreaLuma += LumaTransformSize.GetSize2d(); + if (this.source.IsMonochrome) + { + return; + } + + ObuColorConfig colorConfig = this.picture.Sequence.SequenceHeader.ColorConfig; + int subsamplingX = colorConfig.SubSamplingX ? 1 : 0; + int subsamplingY = colorConfig.SubSamplingY ? 1 : 0; + Point chromaOrigin = new(blockOrigin.X >> subsamplingX, blockOrigin.Y >> subsamplingY); + Av1TransformSize chromaTransformSize = BlockSize.GetMaxUvTransformSize( + colorConfig.SubSamplingX, + colorConfig.SubSamplingY); + + int chromaTransformIndex = this.codedAreaChroma / + Av1EncoderCoefficientBuffer.TransformBlockUnitCoefficientCount; + + ref Av1EncoderTransformBlockState blueState = ref this.blueTransformBlocks[chromaTransformIndex]; + ref Av1EncoderTransformBlockState redState = ref this.redTransformBlocks[chromaTransformIndex]; + this.EncodePlaneBlock( + Av1Plane.U, + chromaOrigin, + chromaTransformSize, + this.blueCoefficients[this.codedAreaChroma..], + ref blueState); + + this.EncodePlaneBlock( + Av1Plane.V, + chromaOrigin, + chromaTransformSize, + this.redCoefficients[this.codedAreaChroma..], + ref redState); + + this.codedAreaChroma += chromaTransformSize.GetSize2d(); + } + + private void EncodePlaneBlock( + Av1Plane plane, + Point blockOrigin, + Av1TransformSize transformSize, + Span coefficients, + ref Av1EncoderTransformBlockState state) + { + Buffer2DRegion sourcePlane = this.source.GetPlane(plane); + Buffer2DRegion reconstructionPlane = this.reconstruction.GetPlane(plane); + int width = transformSize.GetWidth(); + int height = transformSize.GetHeight(); + bool hasLeft = blockOrigin.X > 0; + bool hasAbove = blockOrigin.Y > 0; + ReadOnlySpan above = hasAbove + ? reconstructionPlane.DangerousGetRowSpan(blockOrigin.Y - 1).Slice(blockOrigin.X, width) + : []; + + Span left = TOperator.GetLeftReference(this.blockWorkspace.Residual, height); + if (hasLeft) + { + for (int row = 0; row < height; row++) + { + left[row] = reconstructionPlane.DangerousGetRowSpan(blockOrigin.Y + row)[blockOrigin.X - 1]; + } + } + + // Prediction consumes every gathered reference before residual construction reuses the same workspace bytes. + TOperator.Encode( + this.blockWorkspace, + sourcePlane, + reconstructionPlane, + blockOrigin, + above, + left, + hasLeft, + hasAbove, + coefficients, + transformSize, + this.quantization.QIndex[0], + this.quantization.DeltaQDc[(int)plane], + this.quantization.DeltaQAc[(int)plane], + plane, + this.bitDepth, + ref state); + } + } +} diff --git a/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1IntraSuperblockEncoderTests.cs b/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1IntraSuperblockEncoderTests.cs new file mode 100644 index 0000000000..5a840a1e58 --- /dev/null +++ b/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1IntraSuperblockEncoderTests.cs @@ -0,0 +1,401 @@ +// Copyright (c) Six Labors. +// Licensed under the Six Labors Split License. + +using System.Buffers; +using SixLabors.ImageSharp.Formats.Heif.Av1; +using SixLabors.ImageSharp.Formats.Heif.Av1.Entropy; +using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit; +using SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline; +using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction; +using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling; +using SixLabors.ImageSharp.Formats.Heif.Av1.Transform; +using SixLabors.ImageSharp.Memory; + +namespace SixLabors.ImageSharp.Tests.Formats.Heif.Av1; + +/// +/// Verifies fixed DC intra superblock traversal and reconstruction. +/// +[Trait("Format", "Avif")] +public class Av1IntraSuperblockEncoderTests +{ + [Fact] + public void EncodesClipped128SuperblockInWriterPreorderWithoutAllocation() + { + const int Width = 16; + const int Height = 16; + ObuColorConfig colorConfig = new() + { + IsMonochrome = false, + SubSamplingX = true, + SubSamplingY = true, + BitDepth = Av1BitDepth.EightBit + }; + + using Av1EncoderFrameBuffer source = new( + Configuration.Default, + Width, + Height, + 8, + Av1ColorFormat.Yuv420, + 1, + 1); + + using Av1EncoderFrameBuffer reconstruction = new( + Configuration.Default, + Width, + Height, + 8, + Av1ColorFormat.Yuv420, + 1, + 1); + + FillPlane(source.Frame.CodedView.GetPlane(Av1Plane.Y), 251, 17); + FillPlane(source.Frame.CodedView.GetPlane(Av1Plane.U), 239, 31); + FillPlane(source.Frame.CodedView.GetPlane(Av1Plane.V), 233, 47); + ClearPlane(reconstruction.Luma); + ClearPlane(Assert.IsType>(reconstruction.ChromaBlue)); + ClearPlane(Assert.IsType>(reconstruction.ChromaRed)); + + using Av1EncoderModeInfoBuffer modeInfo = new(Configuration.Default, Width, Height, disallow4x4AllFrames: true); + Av1PictureControlSet picture = CreatePicture(modeInfo, colorConfig, use128x128Superblock: true, qIndex: 73); + using Av1EncoderCoefficientBuffer coefficients = new( + Configuration.Default, + picture.Sequence.SequenceHeader, + Width, + Height); + + using Av1EncoderSuperblockWorkspace superblockWorkspace = new(Configuration.Default); + using Av1EncoderBlockWorkspace blockWorkspace = new(Configuration.Default); + Av1Superblock superblock = new() + { + Workspace = superblockWorkspace, + TileInfo = new Av1TileInfo(0, 0, picture.Parent.FrameHeader), + Index = 0 + }; + + Av1IntraSuperblockEncoder.Encode( + source.Frame, + reconstruction.Frame, + picture, + superblock, + coefficients, + blockWorkspace); + + long before = GC.GetAllocatedBytesForCurrentThread(); + for (int iteration = 0; iteration < 8; iteration++) + { + Av1IntraSuperblockEncoder.Encode( + source.Frame, + reconstruction.Frame, + picture, + superblock, + coefficients, + blockWorkspace); + } + + Assert.Equal(0, GC.GetAllocatedBytesForCurrentThread() - before); + Av1PartitionType[] expectedPartitions = + [ + Av1PartitionType.Split, + Av1PartitionType.Split, + Av1PartitionType.Split, + Av1PartitionType.Split, + Av1PartitionType.None, + Av1PartitionType.None, + Av1PartitionType.None, + Av1PartitionType.None + ]; + + for (int index = 0; index < expectedPartitions.Length; index++) + { + Assert.Equal(expectedPartitions[index], (Av1PartitionType)superblock.CodingUnitPartitionTypes[index]); + } + + for (int index = 0; index < 4; index++) + { + Assert.True(superblock.FinalBlocks[index].HasChroma); + Assert.Equal(73, superblock.FinalBlocks[index].QuantizationIndex); + Assert.Equal(Av1FilterIntraMode.AllFilterIntraModes, superblock.FinalBlocks[index].FilterIntraMode); + } + + Point[] modeInfoPositions = [new(0, 0), new(2, 0), new(0, 2), new(2, 2)]; + foreach (Point position in modeInfoPositions) + { + ref Av1MacroBlockModeInfo block = ref picture.GetMacroBlockModeInfo(position); + Assert.Equal(Av1BlockSize.Block8x8, block.Block.BlockSize); + Assert.Equal(Av1TransformSize.Size8x8, block.Block.TransformSize); + Assert.Equal(Av1PredictionMode.DC, block.Block.Mode); + Assert.Equal(Av1ChromaPredictionMode.DC, block.Block.UvMode); + Assert.False(block.Block.Skip); + } + + Span lumaStates = coefficients.GetTransformBlockSpan(0, Av1Plane.Y); + Span blueStates = coefficients.GetTransformBlockSpan(0, Av1Plane.U); + Span redStates = coefficients.GetTransformBlockSpan(0, Av1Plane.V); + int[] lumaStateIndices = [0, 4, 8, 12]; + for (int index = 0; index < 4; index++) + { + Assert.NotEqual((ushort)0, lumaStates[lumaStateIndices[index]].EndOfBlock); + Assert.Equal(Av1TransformType.DctDct, lumaStates[lumaStateIndices[index]].TransformType); + Assert.NotEqual((ushort)0, blueStates[index].EndOfBlock); + Assert.NotEqual((ushort)0, redStates[index].EndOfBlock); + } + + AssertContainsNonzero(reconstruction.Frame.CodedView.GetPlane(Av1Plane.Y)); + AssertContainsNonzero(reconstruction.Frame.CodedView.GetPlane(Av1Plane.U)); + AssertContainsNonzero(reconstruction.Frame.CodedView.GetPlane(Av1Plane.V)); + + // Edge contexts cover the complete 128x128 superblock because partition updates retain the coded geometry + // even when most of the superblock lies beyond this deliberately clipped frame. + const int ContextUnitCount = 128 >> Av1Constants.ModeInfoSizeLog2; + using Av1NeighborArrayUnit partitions = new( + Configuration.Default, + ContextUnitCount, + ContextUnitCount) + { + GranularityNormalLog2 = Av1Constants.ModeInfoSizeLog2 + }; + + using Av1NeighborArrayUnit lumaContexts = new( + Configuration.Default, + ContextUnitCount, + ContextUnitCount) + { + GranularityNormalLog2 = Av1Constants.ModeInfoSizeLog2 + }; + + using Av1NeighborArrayUnit blueContexts = new( + Configuration.Default, + ContextUnitCount, + ContextUnitCount) + { + GranularityNormalLog2 = Av1Constants.ModeInfoSizeLog2 + }; + + using Av1NeighborArrayUnit redContexts = new( + Configuration.Default, + ContextUnitCount, + ContextUnitCount) + { + GranularityNormalLog2 = Av1Constants.ModeInfoSizeLog2 + }; + + using Av1NeighborArrayUnit transformContexts = new( + Configuration.Default, + ContextUnitCount, + ContextUnitCount) + { + GranularityNormalLog2 = Av1Constants.ModeInfoSizeLog2 + }; + + picture.PartitionContexts = [partitions]; + picture.LuminanceDcSignLevelCoefficientNeighbors = [lumaContexts]; + picture.CbDcSignLevelCoefficientNeighbors = [blueContexts]; + picture.CrDcSignLevelCoefficientNeighbors = [redContexts]; + picture.TransformFunctionContexts = [transformContexts]; + Av1TileWriter.Av1EntropyCodingContext entropyContext = new() + { + MacroBlock = new Av1MacroBlockD { Tile = superblock.TileInfo }, + MacroBlockModeInfo = picture.GetMacroBlockModeInfo(default), + SuperblockOrigin = default + }; + + using Av1SymbolEncoder writer = new(Configuration.Default, 512, 73); + Av1TileWriter.WriteSuperblock( + picture, + entropyContext, + writer, + superblock, + coefficients, + tileIndex: 0); + + using IMemoryOwner encoded = writer.Exit(); + + // The writer must consume exactly the transform areas populated above, proving both traversals stay synchronized. + Assert.Equal(256, entropyContext.CodedAreaSuperblock); + Assert.Equal(64, entropyContext.CodedAreaSuperblockUv); + Assert.NotEqual(0, encoded.GetSpan().Length); + } + + [Fact] + public void PreservesTwelveBitMonochromeReconstructionPrecision() + { + const int Width = 8; + const int Height = 8; + ObuColorConfig colorConfig = new() + { + IsMonochrome = true, + SubSamplingX = true, + SubSamplingY = true, + BitDepth = Av1BitDepth.TwelveBit + }; + + using Av1EncoderFrameBuffer source = new( + Configuration.Default, + Width, + Height, + 12, + Av1ColorFormat.Yuv400, + 0, + 0); + + using Av1EncoderFrameBuffer reconstruction = new( + Configuration.Default, + Width, + Height, + 12, + Av1ColorFormat.Yuv400, + 0, + 0); + + Buffer2DRegion sourcePlane = source.Frame.CodedView.GetPlane(Av1Plane.Y); + for (int y = 0; y < sourcePlane.Height; y++) + { + Span row = sourcePlane.DangerousGetRowSpan(y); + for (int x = 0; x < row.Length; x++) + { + row[x] = (ushort)(3000 + (((x * 71) + (y * 113)) % 1000)); + } + } + + ClearPlane(reconstruction.Luma); + using Av1EncoderModeInfoBuffer modeInfo = new(Configuration.Default, Width, Height, disallow4x4AllFrames: true); + Av1PictureControlSet picture = CreatePicture(modeInfo, colorConfig, use128x128Superblock: false, qIndex: 37); + using Av1EncoderCoefficientBuffer coefficients = new( + Configuration.Default, + picture.Sequence.SequenceHeader, + Width, + Height); + + using Av1EncoderSuperblockWorkspace superblockWorkspace = new(Configuration.Default); + using Av1EncoderBlockWorkspace blockWorkspace = new(Configuration.Default); + Av1Superblock superblock = new() + { + Workspace = superblockWorkspace, + TileInfo = new Av1TileInfo(0, 0, picture.Parent.FrameHeader), + Index = 0 + }; + + Av1IntraSuperblockEncoder.Encode( + source.Frame, + reconstruction.Frame, + picture, + superblock, + coefficients, + blockWorkspace); + + Buffer2DRegion reconstructionPlane = reconstruction.Frame.CodedView.GetPlane(Av1Plane.Y); + ushort maximum = 0; + for (int y = 0; y < reconstructionPlane.Height; y++) + { + foreach (ushort sample in reconstructionPlane.DangerousGetRowSpan(y)) + { + maximum = Math.Max(maximum, sample); + Assert.InRange(sample, (ushort)0, (ushort)4095); + } + } + + Assert.InRange(maximum, (ushort)(byte.MaxValue + 1), (ushort)4095); + Assert.False(superblock.FinalBlocks[0].HasChroma); + Assert.Equal(37, superblock.FinalBlocks[0].QuantizationIndex); + Assert.NotEqual((ushort)0, coefficients.GetTransformBlockSpan(0, Av1Plane.Y)[0].EndOfBlock); + Assert.Equal(0, coefficients.GetPlaneSpan(0, Av1Plane.U).Length); + Assert.Equal(0, coefficients.GetPlaneSpan(0, Av1Plane.V).Length); + } + + private static Av1PictureControlSet CreatePicture( + Av1EncoderModeInfoBuffer modeInfo, + ObuColorConfig colorConfig, + bool use128x128Superblock, + int qIndex) + { + ObuTileGroupHeader tiles = new() + { + TileColumnCount = 1, + TileRowCount = 1 + }; + + tiles.TileColumnStartModeInfo[1] = modeInfo.ModeInfoColumnCount; + tiles.TileRowStartModeInfo[1] = modeInfo.ModeInfoRowCount; + ObuSequenceHeader sequenceHeader = new() + { + Use128x128Superblock = use128x128Superblock, + ColorConfig = colorConfig + }; + + ObuFrameHeader frameHeader = new() + { + ModeInfoColumnCount = modeInfo.ModeInfoColumnCount, + ModeInfoRowCount = modeInfo.ModeInfoRowCount, + TilesInfo = tiles + }; + + frameHeader.QuantizationParameters.BaseQIndex = qIndex; + frameHeader.QuantizationParameters.QIndex.Fill(qIndex); + return new Av1PictureControlSet + { + PartitionContexts = [], + LuminanceDcSignLevelCoefficientNeighbors = [], + CrDcSignLevelCoefficientNeighbors = [], + CbDcSignLevelCoefficientNeighbors = [], + TransformFunctionContexts = [], + Sequence = new Av1SequenceControlSet { SequenceHeader = sequenceHeader }, + Parent = new Av1PictureParentControlSet + { + Common = new Av1EncoderCommon + { + ModeInfoColumnCount = modeInfo.ModeInfoColumnCount, + ModeInfoRowCount = modeInfo.ModeInfoRowCount, + ModeInfoStride = modeInfo.ModeInfoStride, + TilesInfo = tiles, + FrameSize = new ObuFrameSize() + }, + FrameHeader = frameHeader, + PreviousQIndex = [qIndex] + }, + SegmentationNeighborMap = new byte[modeInfo.ModeInfoColumnCount * modeInfo.ModeInfoRowCount], + ModeInfoGrid = modeInfo.Grid, + ModeInfoAllocation = modeInfo.Allocation, + ModeInfoStride = modeInfo.ModeInfoStride, + Disallow4x4AllFrames = modeInfo.Disallow4x4AllFrames, + CdefPreset = [[-1, -1, -1, -1]] + }; + } + + private static void FillPlane(Buffer2DRegion plane, int modulus, int seed) + { + for (int y = 0; y < plane.Height; y++) + { + Span row = plane.DangerousGetRowSpan(y); + for (int x = 0; x < row.Length; x++) + { + row[x] = (byte)(1 + ((seed + (x * 43) + (y * 79)) % modulus)); + } + } + } + + private static void ClearPlane(Buffer2D plane) + where TSample : unmanaged + { + for (int y = 0; y < plane.Height; y++) + { + plane.DangerousGetRowSpan(y).Clear(); + } + } + + private static void AssertContainsNonzero(Buffer2DRegion plane) + where TSample : unmanaged, IEquatable + { + bool containsNonzero = false; + for (int y = 0; y < plane.Height; y++) + { + foreach (TSample sample in plane.DangerousGetRowSpan(y)) + { + containsNonzero |= !sample.Equals(default); + } + } + + Assert.True(containsNonzero); + } +}