From d6a722cc46b9d09c5a3bf7566137cd8d4aac2fc1 Mon Sep 17 00:00:00 2001 From: James Jackson-South Date: Wed, 2 Sep 2026 14:34:57 +1000 Subject: [PATCH] Consolidate AV1 decoder block scratch --- HEIF_IMPLEMENTATION_PLAN.md | 7 + .../Heif/Av1/Transform/Av1BlockDecoder.cs | 128 ++++++++---------- .../Formats/Heif/Av1/Av1FrameBufferTests.cs | 77 ++++++++--- 3 files changed, 120 insertions(+), 92 deletions(-) diff --git a/HEIF_IMPLEMENTATION_PLAN.md b/HEIF_IMPLEMENTATION_PLAN.md index 30a8698a7a..47d3df91f0 100644 --- a/HEIF_IMPLEMENTATION_PLAN.md +++ b/HEIF_IMPLEMENTATION_PLAN.md @@ -721,6 +721,13 @@ Final decoder allocation, lifetime, precision, architecture, and test-validity a boundaries. Frame planes enforce their contiguous single-span invariant before allocation; palette, transform, film-grain, super-resolution, color-conversion, and alpha workspaces remain bounded and allocator owned. No per-block managed allocation remains in reconstruction. +- [x] Block reconstruction now uses one exact-size signed-short owner for inverse quantization, inverse + transform, compound prediction, convolution, and chroma-from-luma scratch. Even-length slices provide + the integer workspaces without another rent. Monochrome reserves no chroma coefficients, and 4:2:0, + 4:2:2, and 4:4:4 reserve two symmetric chroma planes at their coded subsampling. This replaces three + constructor rents and their catch-all rollback path; exact allocation length, coefficient span length, + and exactly-once return pass for all four layouts, with 549 adjacent reconstruction tests passing direct + net11 VSTest in Release. - [x] Valid unsupported tile-list syntax is rejected explicitly. Reserved and metadata OBUs are consumed only after bounded framing and trailing-bit validation. Eight-, ten-, and twelve-bit reconstruction, presentation, alpha, restoration, and film-grain paths retain native precision. diff --git a/src/ImageSharp/Formats/Heif/Av1/Transform/Av1BlockDecoder.cs b/src/ImageSharp/Formats/Heif/Av1/Transform/Av1BlockDecoder.cs index e424f8d32a..fb95ac0a63 100644 --- a/src/ImageSharp/Formats/Heif/Av1/Transform/Av1BlockDecoder.cs +++ b/src/ImageSharp/Formats/Heif/Av1/Transform/Av1BlockDecoder.cs @@ -54,22 +54,27 @@ internal sealed class Av1BlockDecoder : IDisposable private readonly Av1ReferenceFrameStore referenceFrames; /// - /// Owns the reusable raster-order inverse-quantization buffer. + /// Owns all reusable inverse-quantization, transform, and prediction storage. /// - private readonly IMemoryOwner inverseQuantizationOwner; + private readonly IMemoryOwner workspaceOwner; /// - /// Owns the reusable two-dimensional inverse-transform workspace. + /// The inverse-quantization prefix length in signed-short storage elements. /// - private readonly IMemoryOwner transformWorkspaceOwner; + private readonly int inverseQuantizationStorageLength; /// - /// Owns the reusable directional and filter-intra prediction workspace. + /// The inverse-transform workspace offset in signed-short storage elements. /// - private readonly IMemoryOwner predictionScratchOwner; + private readonly int transformWorkspaceOffset; /// - /// The reusable predictor portion of , excluding compound and chroma-from-luma storage. + /// The prediction workspace offset in signed-short storage elements. + /// + private readonly int predictionScratchOffset; + + /// + /// The reusable predictor portion of , excluding compound and chroma-from-luma storage. /// private readonly int predictorWorkingLength; @@ -122,76 +127,56 @@ internal sealed class Av1BlockDecoder : IDisposable // One scratch plane is reused for every transform unit. Its maximum size must cover a complete superblock // across all coded planes, with chroma dimensions reduced independently by their subsampling axes. - int inverseQuantizationSize = ySize + - (this.sequenceHeader.ColorConfig.SubSamplingX ? ySize >> 2 : ySize) + - (this.sequenceHeader.ColorConfig.SubSamplingY ? ySize >> 2 : ySize); - - IMemoryOwner? inverseQuantizationOwner = null; - IMemoryOwner? transformWorkspaceOwner = null; - IMemoryOwner? predictionScratchOwner = null; - try - { - inverseQuantizationOwner = this.frameBuffer.MemoryAllocator.Allocate(inverseQuantizationSize); - transformWorkspaceOwner = this.frameBuffer.MemoryAllocator.Allocate(Av1TransformWorkspace.MaximumLength); - int maximumBlockLength = 1 << sequenceHeader.SuperblockSizeLog2; - int maximumBlockArea = maximumBlockLength * maximumBlockLength; - int predictorWorkingLength = Math.Max( - Av1PredictionDecoder.ScratchLength, - Math.Max( - Av1TranslationalInterPredictor.GetScratchLength(maximumBlockLength, maximumBlockLength), - Av1ScaledInterPredictor.GetMaximumScaledScratchLength(maximumBlockLength, maximumBlockLength))); - - int compoundMaskLength = (maximumBlockArea + 1) >> 1; - int predictorWorkingOffset = (2 * maximumBlockArea) + compoundMaskLength; - int chromaFromLumaOffset = predictorWorkingOffset + predictorWorkingLength; - - // Compound prediction retains both high-precision reference planes plus the full-resolution luma mask. - // Keeping those planes, convolution workspace, and CfL surface in one owner avoids independent managed - // buffers while ensuring the two scratch consumers never overlap. - int predictionScratchLength = chromaFromLumaOffset + Av1ChromaFromLumaContext.BufferLength; - predictionScratchOwner = this.frameBuffer.MemoryAllocator.Allocate(predictionScratchLength); - - this.inverseQuantizationOwner = inverseQuantizationOwner; - this.transformWorkspaceOwner = transformWorkspaceOwner; - this.predictionScratchOwner = predictionScratchOwner; - this.predictorWorkingLength = predictorWorkingLength; - this.predictionDecoder = new( - sequenceHeader, - frameHeader, - predictionScratchOwner.Memory.Slice(predictorWorkingOffset, predictorWorkingLength), - paletteColorIndexMaps); - this.isLoopFilterEnabled = frameHeader.LoopFilterParameters.FilterLevel[0] != 0 || - frameHeader.LoopFilterParameters.FilterLevel[1] != 0; - - this.chromaFromLumaContext = new( - sequenceHeader.ColorConfig, - predictionScratchOwner.Memory.Slice(chromaFromLumaOffset, Av1ChromaFromLumaContext.BufferLength)); - } - catch - { - // A constructor that does not return transfers no ownership to its caller. Unwind successful rents in - // reverse order so allocator diagnostics and pooled buffers remain balanced after any later allocation. - predictionScratchOwner?.Dispose(); - transformWorkspaceOwner?.Dispose(); - inverseQuantizationOwner?.Dispose(); - throw; - } + ObuColorConfig colorConfig = this.sequenceHeader.ColorConfig; + int chromaSubsampling = (colorConfig.SubSamplingX ? 1 : 0) + (colorConfig.SubSamplingY ? 1 : 0); + int chromaSize = ySize >> chromaSubsampling; + int inverseQuantizationSize = colorConfig.IsMonochrome ? ySize : ySize + (2 * chromaSize); + int maximumBlockLength = 1 << sequenceHeader.SuperblockSizeLog2; + int maximumBlockArea = maximumBlockLength * maximumBlockLength; + int predictorWorkingLength = Math.Max( + Av1PredictionDecoder.ScratchLength, + Math.Max( + Av1TranslationalInterPredictor.GetScratchLength(maximumBlockLength, maximumBlockLength), + Av1ScaledInterPredictor.GetMaximumScaledScratchLength(maximumBlockLength, maximumBlockLength))); + + int compoundMaskLength = (maximumBlockArea + 1) >> 1; + int predictorWorkingOffset = (2 * maximumBlockArea) + compoundMaskLength; + int chromaFromLumaOffset = predictorWorkingOffset + predictorWorkingLength; + int predictionScratchLength = chromaFromLumaOffset + Av1ChromaFromLumaContext.BufferLength; + this.inverseQuantizationStorageLength = inverseQuantizationSize * 2; + this.transformWorkspaceOffset = this.inverseQuantizationStorageLength; + this.predictionScratchOffset = this.transformWorkspaceOffset + (Av1TransformWorkspace.MaximumLength * 2); + + // Integer workspaces occupy even signed-short slices so one allocator owner can retain the complete block + // lifetime while prediction still receives the Memory contract needed by its reusable context. + this.workspaceOwner = this.frameBuffer.MemoryAllocator.Allocate( + this.predictionScratchOffset + predictionScratchLength); + + Memory predictionScratch = this.workspaceOwner.Memory[this.predictionScratchOffset..]; + this.predictorWorkingLength = predictorWorkingLength; + this.predictionDecoder = new( + sequenceHeader, + frameHeader, + predictionScratch.Slice(predictorWorkingOffset, predictorWorkingLength), + paletteColorIndexMaps); + this.isLoopFilterEnabled = frameHeader.LoopFilterParameters.FilterLevel[0] != 0 || + frameHeader.LoopFilterParameters.FilterLevel[1] != 0; + + this.chromaFromLumaContext = new( + sequenceHeader.ColorConfig, + predictionScratch.Slice(chromaFromLumaOffset, Av1ChromaFromLumaContext.BufferLength)); } /// /// Gets the reusable raster-order coefficient buffer populated by inverse quantization. /// - public Span CurrentInverseQuantizationCoefficients => this.inverseQuantizationOwner.Memory.Span; + public Span CurrentInverseQuantizationCoefficients + => MemoryMarshal.Cast(this.workspaceOwner.Memory.Span[..this.inverseQuantizationStorageLength]); /// /// Releases the pooled reconstruction workspaces owned by this decoder. /// - public void Dispose() - { - this.predictionScratchOwner.Dispose(); - this.transformWorkspaceOwner.Dispose(); - this.inverseQuantizationOwner.Dispose(); - } + public void Dispose() => this.workspaceOwner.Dispose(); /// /// Resets the per-plane packed coefficient cursors before reconstructing a superblock. @@ -216,7 +201,10 @@ internal sealed class Av1BlockDecoder : IDisposable /// The tile boundaries used to determine neighbor availability. public void DecodeBlock(Av1BlockModeInfo modeInfo, Point modeInfoPosition, Av1BlockSize blockSize, Av1SuperblockInfo superblockInfo, Av1TileInfo tileInfo) { - Span transformWorkspace = this.transformWorkspaceOwner.Memory.Span; + Span transformWorkspace = MemoryMarshal.Cast( + this.workspaceOwner.Memory.Span.Slice( + this.transformWorkspaceOffset, + Av1TransformWorkspace.MaximumLength * 2)); ObuColorConfig colorConfig = this.sequenceHeader.ColorConfig; Av1TransformType transformType; @@ -351,7 +339,7 @@ internal sealed class Av1BlockDecoder : IDisposable int maximumBlockLength = 1 << this.sequenceHeader.SuperblockSizeLog2; int maximumBlockArea = maximumBlockLength * maximumBlockLength; int compoundMaskStorageLength = (maximumBlockArea + 1) >> 1; - Span predictionStorage = this.predictionScratchOwner.Memory.Span; + Span predictionStorage = this.workspaceOwner.Memory.Span[this.predictionScratchOffset..]; Span secondPredictionStorage = predictionStorage[..maximumBlockArea]; Span firstCompoundPrediction = MemoryMarshal.Cast( predictionStorage.Slice(maximumBlockArea, maximumBlockArea)); diff --git a/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1FrameBufferTests.cs b/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1FrameBufferTests.cs index a3e0c53451..f895399719 100644 --- a/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1FrameBufferTests.cs +++ b/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1FrameBufferTests.cs @@ -5,6 +5,9 @@ using SixLabors.ImageSharp.Formats.Heif.Av1; using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit; using SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline.LoopFilter; using SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline.Quantizers; +using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction; +using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.ChromaFromLuma; +using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.Inter; using SixLabors.ImageSharp.Formats.Heif.Av1.ReferenceFrames; using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling; using SixLabors.ImageSharp.Formats.Heif.Av1.Transform; @@ -115,16 +118,20 @@ public class Av1FrameBufferTests } /// - /// Verifies that each later block-decoder workspace failure releases every workspace rented earlier. + /// Verifies that block reconstruction uses one exact-size owner across monochrome and chroma plane layouts. /// [Theory] - [InlineData(3, 1)] - [InlineData(4, 2)] - public void BlockDecoderConstructorFailureReleasesEarlierWorkspaces( - int failureAllocationNumber, - int successfulWorkspaceCount) + [InlineData(true, false, false, 4096)] + [InlineData(false, true, true, 6144)] + [InlineData(false, true, false, 8192)] + [InlineData(false, false, false, 12288)] + public void BlockDecoderUsesOneContiguousWorkspaceOwner( + bool isMonochrome, + bool subsamplingX, + bool subsamplingY, + int expectedInverseQuantizationSize) { - FailingTestMemoryAllocator allocator = new(failureAllocationNumber); + TestMemoryAllocator allocator = new(); Configuration configuration = Configuration.Default.Clone(); configuration.MemoryAllocator = allocator; ObuSequenceHeader sequenceHeader = new() @@ -134,12 +141,20 @@ public class Av1FrameBufferTests Use128x128Superblock = false, ColorConfig = new ObuColorConfig { - IsMonochrome = true, + IsMonochrome = isMonochrome, + SubSamplingX = subsamplingX, + SubSamplingY = subsamplingY, BitDepth = Av1BitDepth.EightBit } }; - using Av1FrameBuffer frameBuffer = new(configuration, sequenceHeader, Av1ColorFormat.Yuv400, false); + Av1ColorFormat colorFormat = isMonochrome + ? Av1ColorFormat.Yuv400 + : subsamplingX + ? subsamplingY ? Av1ColorFormat.Yuv420 : Av1ColorFormat.Yuv422 + : Av1ColorFormat.Yuv444; + + using Av1FrameBuffer frameBuffer = new(configuration, sequenceHeader, colorFormat, false); ObuFrameHeader frameHeader = new() { ModeInfoColumnCount = 16, @@ -152,28 +167,46 @@ public class Av1FrameBufferTests Av1InverseQuantizer inverseQuantizer = new(sequenceHeader, frameHeader); using Av1ReferenceFrameStore referenceFrames = new(); - // The frame's luma plane is allocation attempt one. Resetting only the logs preserves that counter while - // isolating the block-decoder owners that must be returned when a later workspace rent fails. + // Reset the frame-plane logs so the following assertions describe only the block decoder's scratch owner. allocator.EnableNonThreadSafeLogging(); - Assert.Throws( - () => new Av1BlockDecoder( + int maximumBlockLength = 1 << sequenceHeader.SuperblockSizeLog2; + int maximumBlockArea = maximumBlockLength * maximumBlockLength; + int predictorWorkingLength = Math.Max( + Av1PredictionDecoder.ScratchLength, + Math.Max( + Av1TranslationalInterPredictor.GetScratchLength(maximumBlockLength, maximumBlockLength), + Av1ScaledInterPredictor.GetMaximumScaledScratchLength(maximumBlockLength, maximumBlockLength))); + + int predictionScratchLength = + (2 * maximumBlockArea) + + ((maximumBlockArea + 1) >> 1) + + predictorWorkingLength + + Av1ChromaFromLumaContext.BufferLength; + + int expectedWorkspaceLength = + (expectedInverseQuantizationSize * 2) + + (Av1TransformWorkspace.MaximumLength * 2) + + predictionScratchLength; + + TestMemoryAllocator.AllocationRequest workspaceAllocation; + using (Av1BlockDecoder blockDecoder = new( sequenceHeader, frameHeader, frameBuffer, loopFilterContext, inverseQuantizer, - referenceFrames)); - - Assert.Equal(failureAllocationNumber, allocator.AllocationAttemptCount); - Assert.Equal(successfulWorkspaceCount, allocator.AllocationLog.Count); - Assert.Equal(successfulWorkspaceCount, allocator.ReturnLog.Count); - foreach (TestMemoryAllocator.AllocationRequest allocation in allocator.AllocationLog) + referenceFrames)) { - Assert.Contains( - allocator.ReturnLog, - returned => returned.HashCodeOfBuffer == allocation.HashCodeOfBuffer); + workspaceAllocation = Assert.Single(allocator.AllocationLog); + Assert.Empty(allocator.ReturnLog); + Assert.Equal(typeof(short), workspaceAllocation.ElementType); + Assert.Equal(expectedWorkspaceLength, workspaceAllocation.Length); + Assert.Equal(expectedInverseQuantizationSize, blockDecoder.CurrentInverseQuantizationCoefficients.Length); } + + TestMemoryAllocator.ReturnRequest returned = Assert.Single(allocator.ReturnLog); + Assert.Equal(workspaceAllocation.AllocationId, returned.AllocationId); } ///