Browse Source

Add DCT, transpose, and performance improvements

- Add a slice & assert to JxlHuffmanDecoder alphabetSize to allocate at most 256 items
- Use [0, 0] instead of stackalloc[2] followed by Clear() in JxlAnsReader
- Add assert to Butteraugli ComputeKernel method & use float for Butteraugli Wmul & use InlineArray
- Add transpose.
    - Note: transpose is scalar, it doesn't support SIMD yet
- Add shared constants & file signature
- Improve while loop in JxlImageOperations.Mirror
- Floating-point Discrete Cosine Transform (1D and 2D)
- Add an inline array of 2 items

Source files implemented from libjxl with this commit:
- dct-inl.h
- dct_block-inl.h
- transpose-inl.h
pull/3153/head
winscripter 3 weeks ago
parent
commit
27a848cd25
  1. 9
      src/ImageSharp/Formats/Jxl/InlineArrays.cs
  2. 20
      src/ImageSharp/Formats/Jxl/Processing/Butteraugli/Butteraugli.cs
  3. 3
      src/ImageSharp/Formats/Jxl/Processing/Decoder/JxlAnsReader.cs
  4. 5
      src/ImageSharp/Formats/Jxl/Processing/Decoder/JxlHuffmanDecoder.cs
  5. 305
      src/ImageSharp/Formats/Jxl/Processing/JxlDct.cs
  6. 52
      src/ImageSharp/Formats/Jxl/Processing/JxlDctOutput.cs
  7. 27
      src/ImageSharp/Formats/Jxl/Processing/JxlDctScales.cs
  8. 55
      src/ImageSharp/Formats/Jxl/Processing/JxlDctSource.cs
  9. 8
      src/ImageSharp/Formats/Jxl/Processing/JxlImageOperations.cs
  10. 30
      src/ImageSharp/Formats/Jxl/Processing/JxlShared.cs
  11. 22
      src/ImageSharp/Formats/Jxl/Processing/JxlTranspose.cs

9
src/ImageSharp/Formats/Jxl/InlineArrays.cs

@ -7,6 +7,15 @@ using System.Runtime.CompilerServices;
namespace SixLabors.ImageSharp.Formats.Jxl;
/// <summary>
/// Used by Butteraugli
/// </summary>
[InlineArray(2)]
internal struct InlineArray2<T>
{
private T first;
}
[InlineArray(3)]
internal struct InlineArray3<T>
{

20
src/ImageSharp/Formats/Jxl/Processing/Butteraugli/Butteraugli.cs

@ -42,11 +42,11 @@ internal static class Butteraugli
private const float GlobalScale =
1.0f / InternalGoodQualityThreshold;
public static ReadOnlySpan<double> Wmul =>
public static ReadOnlySpan<float> Wmul =>
[
400.0, 1.50815703118, 0,
2150.0, 10.6195433239, 16.2176043152,
29.2353797994, 0.844626970982, 0.703646627719,
400.0f, 1.50815703118f, 0f,
2150.0f, 10.6195433239f, 16.2176043152f,
29.2353797994f, 0.844626970982f, 0.703646627719f,
];
public static ReadOnlySpan<float> ComputeKernel(float sigma)
@ -55,7 +55,11 @@ internal static class Butteraugli
float scaler = -1.0f / (2.0f * sigma * sigma);
int diff = Math.Max(1, (int)(m * MathF.Abs(sigma)));
// Use new because there's only up to 3 elements
// If sigma is very large we should not return a 'new float[]' allocation.
// This guard is temporary so we can verify the range of the number of elements.
// TODO: remove guard if the value doesn't exceed the limit for many JXL files
DebugGuard.MustBeLessThanOrEqualTo(sigma, 32f, nameof(sigma));
float[] kernel = new float[(2 * diff) + 1];
for (int i = -diff; i <= diff; i++)
@ -536,7 +540,7 @@ internal static class Butteraugli
Configuration configuration,
in ButteraugliParameters parameters,
JxlImage3F mf,
JxlImageF[] hf,
ref InlineArray2<JxlImageF> hf,
BlurTemp blurTemp)
{
const float sigmaHf = 3.22489901262f;
@ -2095,8 +2099,8 @@ internal static class Butteraugli
}
}
JxlImageF[] hf0 = new JxlImageF[2];
JxlImageF[] hf1 = new JxlImageF[2];
InlineArray2<JxlImageF> hf0 = default;
InlineArray2<JxlImageF> hf1 = default;
if (!SeparateMfAndHf(parameters, image0, hf0, blurTemp))
{

3
src/ImageSharp/Formats/Jxl/Processing/Decoder/JxlAnsReader.cs

@ -103,8 +103,7 @@ internal static class JxlAnsReader
if (isSimpleCode)
{
Span<uint> symbols = stackalloc uint[2];
symbols.Clear();
Span<uint> symbols = [0, 0];
uint maxSymbol = 0u;
uint symCount = reader.ReadBits32(1u) + 1u;

5
src/ImageSharp/Formats/Jxl/Processing/Decoder/JxlHuffmanDecoder.cs

@ -269,8 +269,9 @@ internal sealed class JxlHuffmanDecoder
return ReadSimpleCode(alphabetSize, br, this.Table);
}
// The alphabet size is at most 256
Span<byte> codeLengths = stackalloc byte[alphabetSize];
DebugGuard.MustBeLessThanOrEqualTo(alphabetSize, 256, nameof(alphabetSize));
Span<byte> codeLengths = stackalloc byte[256].Slice(0, alphabetSize);
codeLengths.Clear(); // Zero-initialized in reference software
Span<byte> codeLengthCodeLengths = stackalloc byte[CodeLengthCodes];

305
src/ImageSharp/Formats/Jxl/Processing/JxlDct.cs

@ -0,0 +1,305 @@
// Copyright (c) Six Labors.
// Licensed under the Six Labors Split License.
using System.Numerics;
using System.Runtime.CompilerServices;
namespace SixLabors.ImageSharp.Formats.Jxl.Processing;
/// <summary>
/// Discrete Cosine Transform with SIMD support.
/// </summary>
internal static class JxlDct
{
/// <summary>
/// Creates a new coefficient bundle.
/// </summary>
/// <param name="n">Number of items.</param>
/// <param name="sz">Coefficient size.</param>
/// <returns>A new coefficient bundle.</returns>
public static CoefficientBundle CoeffBundle(int n, int sz) => new(n, sz);
public static void Dct1DCore(int n, int sz, Span<float> mem, Span<float> tmp)
{
if (n == 2)
{
Vector<float> in1 = new(mem);
Vector<float> in2 = new(mem[sz..]);
(in1 + in2).CopyTo(mem);
(in1 - in2).CopyTo(mem[sz..]);
}
else
{
CoefficientBundle cb = CoeffBundle(n / 2, sz);
cb.AddReverse(mem, mem[(n / 2 * sz)..], tmp);
Dct1DCore(n / 2, sz, tmp, tmp[(n * sz)..]);
cb.SubReverse(mem, mem[(n / 2 * sz)..], tmp[(n / 2 * sz)..]);
cb.Multiply(tmp);
Dct1DCore(n / 2, sz, tmp[(n / 2 * sz)..], tmp[(n * sz)..]);
cb.B(tmp[(n / 2 * sz)..]);
CoeffBundle(n, sz).InverseEvenOdd(tmp, mem);
}
}
public static void InverseDct1DCore(int n, int sz, Span<float> from, int fromStride, Span<float> to, int toStride, Span<float> tmp)
{
if (n == 1)
{
from.CopyTo(to);
}
else if (n == 2)
{
Vector<float> in1 = new(from);
Vector<float> in2 = new(from[fromStride..]);
(in1 + in2).CopyTo(to);
(in1 + in2).CopyTo(to[toStride..]);
}
else
{
CoefficientBundle cbDiv2 = CoeffBundle(n / 2, sz);
CoefficientBundle cb = CoeffBundle(n, sz);
cb.ForwardEvenOdd(from, fromStride, tmp);
InverseDct1DCore(n / 2, sz, tmp, sz, tmp, sz, tmp[(n * sz)..]);
cbDiv2.BTranspose(tmp[((n / 2) * sz)..]);
InverseDct1DCore(n / 2, sz, tmp[((n / 2) * sz)..], sz, tmp[((n / 2) * sz)..], sz, tmp[(n * sz)..]);
cb.MultiplyAndAdd(tmp, to, toStride);
}
}
public static void Dct1DWrapper(int n, int m, bool fit, JxlDctSource from, JxlDctOutput to, int mp, Span<float> tmp)
{
CoefficientBundle cb = CoeffBundle(n, m);
for (int i = 0; i < mp; i += m)
{
cb.LoadFromBlock(from, i, tmp);
Dct1DCore(n, m, tmp, tmp[(n * m)..]);
cb.StoreToBlockAndScale(tmp, ref to, i);
if (fit)
{
return;
}
}
}
public static void InverseDct1DWrapper(int n, int m, bool fit, JxlDctSource from, JxlDctOutput to, int mp, Span<float> tmp)
{
for (int i = 0; i < mp; i += m)
{
InverseDct1DCore(n, m, from.Address(0, i), from.Stride, to.Address(0, i), to.Stride, tmp);
if (fit)
{
return;
}
}
}
public static void Dct1DCapped(int n, int m, int l, JxlDctSource from, JxlDctOutput to, Span<float> tmp)
{
bool fit = m <= l;
Dct1DWrapper(n, m, fit, from, to, m, tmp);
}
public static void InverseDct1DCapped(int n, int m, int l, JxlDctSource from, JxlDctOutput to, Span<float> tmp)
{
bool fit = m <= l;
InverseDct1DWrapper(n, m, fit, from, to, m, tmp);
}
public static void Dct1D(int n, int m, JxlDctSource from, JxlDctOutput to, Span<float> tmp)
{
int lanes = Vector<float>.Count;
Dct1DCapped(n, m, lanes, from, to, tmp);
}
public static void InverseDct1D(int n, int m, JxlDctSource source, JxlDctOutput output, Span<float> tmp)
{
int lanes = Vector<float>.Count;
InverseDct1DCapped(n, m, lanes, source, output, tmp);
}
public static void ComputeScaledDct(int rows, int columns, JxlDctSource from, Span<float> to, Span<float> scratchSpace)
{
Span<float> block = scratchSpace;
Span<float> tmp = scratchSpace[(rows * columns)..];
if (rows < columns)
{
Dct1D(rows, columns, from, new JxlDctOutput(block, columns), tmp);
JxlTranspose.Transpose(rows, columns, new JxlDctSource(block, columns), new JxlDctOutput(to, rows));
Dct1D(columns, rows, new JxlDctSource(to, rows), new JxlDctOutput(block, rows), tmp);
JxlTranspose.Transpose(columns, rows, new JxlDctSource(block, rows), new JxlDctOutput(to, columns));
}
else
{
Dct1D(rows, columns, from, new JxlDctOutput(to, columns), tmp);
JxlTranspose.Transpose(rows, columns, new JxlDctSource(to, columns), new JxlDctOutput(block, rows));
Dct1D(columns, rows, new JxlDctSource(block, rows), new JxlDctOutput(to, rows), tmp);
}
}
public static void ComputeScaledInverseDct(int rows, int columns, Span<float> from, JxlDctOutput to, Span<float> scratchSpace)
{
Span<float> block = scratchSpace;
Span<float> tmp = scratchSpace[(rows * columns)..];
if (rows < columns)
{
JxlTranspose.Transpose(rows, columns, new JxlDctSource(from, columns), new JxlDctOutput(block, rows));
InverseDct1D(columns, rows, new JxlDctSource(block, rows), new JxlDctOutput(from, rows), tmp);
JxlTranspose.Transpose(columns, rows, new JxlDctSource(from, rows), new JxlDctOutput(block, columns));
InverseDct1D(rows, columns, new JxlDctSource(block, columns), to, tmp);
}
else
{
InverseDct1D(columns, rows, new JxlDctSource(from, rows), new JxlDctOutput(block, rows), tmp);
JxlTranspose.Transpose(columns, rows, new JxlDctSource(block, rows), new JxlDctOutput(from, columns));
InverseDct1D(rows, columns, new JxlDctSource(from, columns), to, tmp);
}
}
/// <summary>
/// Core methods for the Discrete Cosine Transform (DCT).
/// </summary>
public readonly struct CoefficientBundle(int n, int sz)
{
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void AddReverse(Span<float> aIn1, Span<float> aIn2, Span<float> aOut)
{
for (int i = 0; i < n; i++)
{
Vector<float> in1 = new(aIn1[(i * sz)..]);
Vector<float> in2 = new(aIn2[((n - i - 1) * sz)..]);
(in1 + in2).CopyTo(aOut[(i * sz)..]);
}
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void SubReverse(Span<float> aIn1, Span<float> aIn2, Span<float> aOut)
{
for (int i = 0; i < n; i++)
{
Vector<float> in1 = new(aIn1[(i * sz)..]);
Vector<float> in2 = new(aIn2[((n - i - 1) * sz)..]);
(in1 - in2).CopyTo(aOut[(i * sz)..]);
}
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void B(Span<float> coeff)
{
Vector<float> sqrt2 = new(JxlDctScales.Sqrt2);
Vector<float> in10 = new(coeff);
Vector<float> in20 = new(coeff[sz..]);
((in10 * sqrt2) + in20).CopyTo(coeff);
for (int i = 1; i + 1 < n; i++)
{
Vector<float> in1 = new(coeff[(i * sz)..]);
Vector<float> in2 = new(coeff[((i + 1) * sz)..]);
(in1 + in2).CopyTo(coeff[(i * sz)..]);
}
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void BTranspose(Span<float> coeff)
{
for (int i = n - 1; i > 0; i--)
{
Vector<float> in1 = new(coeff[(i * sz)..]);
Vector<float> in2 = new(coeff[((i - 1) * sz)..]);
(in1 + in2).CopyTo(coeff[(i * sz)..]);
}
Vector<float> sqrt2 = new(JxlDctScales.Sqrt2);
Vector<float> in1x = new(coeff);
(in1x * sqrt2).CopyTo(coeff);
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void InverseEvenOdd(Span<float> aIn, Span<float> aOut)
{
for (int i = 0; i < n / 2; i++)
{
new Vector<float>(aIn[(i * sz)..]).CopyTo(aOut[((2 * i) * sz)..]);
}
for (int i = n / 2; i < n; i++)
{
new Vector<float>(aIn[(i * sz)..]).CopyTo(aOut[(((2 * (i - (n / 2))) + 1) * sz)..]);
}
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void ForwardEvenOdd(Span<float> aIn, int aInStride, Span<float> aOut)
{
for (int i = 0; i < n / 2; i++)
{
new Vector<float>(aIn[(2 * i * aInStride)..]).CopyTo(aOut[(i * sz)..]);
}
for (int i = n / 2; i < n; i++)
{
new Vector<float>(aIn[(((2 * (i - (n / 2))) + 1) * aInStride)..]).CopyTo(aOut[(i * sz)..]);
}
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void Multiply(Span<float> coeff)
{
ReadOnlySpan<float> multipliers = JxlDctScales.GetMultipliers(n);
for (int i = 0; i < n / 2; i++)
{
Vector<float> in1 = new(coeff[(((n / 2) + i) * sz)..]);
Vector<float> mul = new(multipliers[i]);
(in1 * mul).CopyTo(coeff[((n / (2 + i)) * sz)..]);
}
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void MultiplyAndAdd(Span<float> coeff, Span<float> output, int outStride)
{
ReadOnlySpan<float> multipliers = JxlDctScales.GetMultipliers(n);
for (int i = 0; i < n / 2; i++)
{
Vector<float> mul = new(multipliers[i]);
Vector<float> in1 = new(coeff[(i * sz)..]);
Vector<float> in2 = new(coeff[((n / (2 + i)) * sz)..]);
Vector<float> out1 = (mul * in2) * in1;
Vector<float> out2 = -(mul * in2) + in1;
out1.CopyTo(output[(i * outStride)..]);
out2.CopyTo(output[((n - i - 1) * outStride)..]);
}
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void LoadFromBlock(in JxlDctSource input, int offset, Span<float> coeff)
{
for (int i = 0; i < n; i++)
{
input.LoadPart(i, offset).CopyTo(coeff[(i * sz)..]);
}
}
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void StoreToBlockAndScale(Span<float> coeff, ref JxlDctOutput output, int offset)
{
Vector<float> mul = new(1.0f / n);
for (int i = 0; i < n; i++)
{
output.StorePart(mul * new Vector<float>(coeff[(i * sz)..]), i, offset);
}
}
}
}

52
src/ImageSharp/Formats/Jxl/Processing/JxlDctOutput.cs

@ -0,0 +1,52 @@
// Copyright (c) Six Labors.
// Licensed under the Six Labors Split License.
using System.Numerics;
using System.Runtime.CompilerServices;
namespace SixLabors.ImageSharp.Formats.Jxl.Processing;
/// <summary>
/// Output DCT block.
/// </summary>
internal ref struct JxlDctOutput(Span<float> data, int stride)
{
/// <summary>
/// Raw block data.
/// </summary>
public Span<float> Data = data;
/// <summary>
/// Stride size.
/// </summary>
public readonly int Stride = stride;
/// <summary>
/// Returns the span to the start of a row and offset.
/// </summary>
/// <param name="row">The row index.</param>
/// <param name="i">The offset.</param>
/// <returns>
/// Span for that row &amp; offset.
/// </returns>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public readonly Span<float> Address(int row, int i) => this.Data[((row * this.Stride) + i)..];
/// <summary>
/// Writes a single value to the block at the row and offset.
/// </summary>
/// <param name="value">The value to write.</param>
/// <param name="row">The row index.</param>
/// <param name="i">The offset.</param>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void Write(float value, int row, int i) => this.Data[(row * this.Stride) + i] = value;
/// <summary>
/// Stores the vector into the data at the specified row and offset.
/// </summary>
/// <param name="value">The vector to write.</param>
/// <param name="row">The row index.</param>
/// <param name="index">The offset.</param>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public readonly void StorePart(Vector<float> value, int row, int index) => value.CopyTo(this.Address(row, index));
}

27
src/ImageSharp/Formats/Jxl/Processing/JxlDctScales.cs

@ -9,6 +9,16 @@ namespace SixLabors.ImageSharp.Formats.Jxl.Processing;
/// </summary>
internal static class JxlDctScales
{
/// <summary>
/// Square root of 2.
/// </summary>
public const float Sqrt2 = 1.41421356237f;
/// <summary>
/// Square root of 0.5.
/// </summary>
public const float Sqrt05 = 0.70710678118f;
/// <summary>
/// Gets 8x1 DCT resample scales.
/// </summary>
@ -358,4 +368,21 @@ internal static class JxlDctScales
9.058751453879703f, 11.644627325175037f, 16.300023088031555f,
27.163977662448232f, 81.48784219222516f,
];
/// <summary>
/// Returns DCT multipliers for size <paramref name="n"/>.
/// </summary>
/// <param name="n">The multiplier size</param>
/// <returns>DCT multipliers for the given size.</returns>
public static ReadOnlySpan<float> GetMultipliers(int n) => n switch
{
4 => Multipliers4,
8 => Multipliers8,
16 => Multipliers16,
32 => Multipliers32,
64 => Multipliers64,
128 => Multipliers128,
256 => Multipliers256,
_ => []
};
}

55
src/ImageSharp/Formats/Jxl/Processing/JxlDctSource.cs

@ -0,0 +1,55 @@
// Copyright (c) Six Labors.
// Licensed under the Six Labors Split License.
using System.Numerics;
using System.Runtime.CompilerServices;
namespace SixLabors.ImageSharp.Formats.Jxl.Processing;
/// <summary>
/// Source DCT block.
/// </summary>
internal readonly ref struct JxlDctSource(Span<float> data, int stride)
{
/// <summary>
/// Raw block data.
/// </summary>
public readonly Span<float> Data = data;
/// <summary>
/// Stride size.
/// </summary>
public readonly int Stride = stride;
/// <summary>
/// Returns the span to the start of a row and offset.
/// </summary>
/// <param name="row">The row index.</param>
/// <param name="i">The offset.</param>
/// <returns>
/// Span for that row &amp; offset.
/// </returns>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public Span<float> Address(int row, int i) => this.Data[((row * this.Stride) + i)..];
/// <summary>
/// Returns the coefficient at the row and offset.
/// </summary>
/// <param name="row">The row index.</param>
/// <param name="i">The offset.</param>
/// <returns>
/// Coefficient at that row and offset.
/// </returns>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public float Read(int row, int i) => this.Data[(row * this.Stride) + i];
/// <summary>
/// Loads a vector at the specified row and offset.
/// </summary>
/// <param name="row">The row index.</param>
/// <param name="i">The offset.</param>
/// <returns>
/// Vector at that row and offset.
/// </returns>
public Vector<float> LoadPart(int row, int i) => new(this.Address(row, i));
}

8
src/ImageSharp/Formats/Jxl/Processing/JxlImageOperations.cs

@ -346,16 +346,20 @@ internal static class JxlImageOperations
{
DebugGuard.MustBeGreaterThan(xSize, 0, nameof(xSize));
while (x < 0 || x >= xSize)
while (true)
{
if (x < 0)
{
x = -x - 1;
}
else
else if (x >= xSize)
{
x = (2 * xSize) - 1 - x;
}
else
{
break;
}
}
return (int)x;

30
src/ImageSharp/Formats/Jxl/Processing/JxlShared.cs

@ -0,0 +1,30 @@
// Copyright (c) Six Labors.
// Licensed under the Six Labors Split License.
namespace SixLabors.ImageSharp.Formats.Jxl.Processing;
/// <summary>
/// Shared JPEG XL constants
/// </summary>
internal static class JxlShared
{
/// <summary>
/// Maximum number of passes in an image.
/// </summary>
public const int MaximumNumberOfPasses = 11;
/// <summary>
/// Maximum number of reference frames.
/// </summary>
public const int MaximumNumberOfReferenceFrames = 4;
/// <summary>
/// Gets the 12-byte signature (a.k.a. magic) for JPEG XL files.
/// </summary>
public static ReadOnlySpan<byte> SignatureBox =>
[
0x00, 0x00, 0x00, 0x0C,
(byte)'J', (byte)'X', (byte)'L', (byte)' ',
0x0D, 0x0A, 0x87, 0x0A
];
}

22
src/ImageSharp/Formats/Jxl/Processing/JxlTranspose.cs

@ -0,0 +1,22 @@
// Copyright (c) Six Labors.
// Licensed under the Six Labors Split License.
namespace SixLabors.ImageSharp.Formats.Jxl.Processing;
/// <summary>
/// Performs transpose on JPEG XL DCT blocks.
/// </summary>
internal static class JxlTranspose
{
// TODO: SIMD
public static void Transpose(int r, int c, JxlDctSource from, JxlDctOutput to)
{
for (int n = 0; n < r; n++)
{
for (int m = 0; m < c; m++)
{
to.Write(from.Read(n, m), m, n);
}
}
}
}
Loading…
Cancel
Save