Coverage Report

Created: 2025-06-13 07:37

/src/libjxl/lib/jxl/dec_noise.cc
Line
Count
Source
1
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
2
//
3
// Use of this source code is governed by a BSD-style
4
// license that can be found in the LICENSE file.
5
6
#include "lib/jxl/dec_noise.h"
7
8
#include <cstdint>
9
#include <cstdlib>
10
#include <cstring>
11
#include <utility>
12
13
#include "lib/jxl/base/status.h"
14
#include "lib/jxl/dec_bit_reader.h"
15
#include "lib/jxl/dec_cache.h"
16
#include "lib/jxl/frame_header.h"
17
#include "lib/jxl/noise.h"
18
#include "lib/jxl/render_pipeline/render_pipeline.h"
19
20
#undef HWY_TARGET_INCLUDE
21
#define HWY_TARGET_INCLUDE "lib/jxl/dec_noise.cc"
22
#include <hwy/foreach_target.h>
23
#include <hwy/highway.h>
24
25
#include "lib/jxl/base/compiler_specific.h"
26
#include "lib/jxl/base/rect.h"
27
#include "lib/jxl/frame_dimensions.h"
28
#include "lib/jxl/image.h"
29
#include "lib/jxl/xorshift128plus-inl.h"
30
31
HWY_BEFORE_NAMESPACE();
32
namespace jxl {
33
namespace HWY_NAMESPACE {
34
35
// These templates are not found via ADL.
36
using hwy::HWY_NAMESPACE::Or;
37
using hwy::HWY_NAMESPACE::Rebind;
38
using hwy::HWY_NAMESPACE::ShiftRight;
39
using hwy::HWY_NAMESPACE::Vec;
40
41
using D = HWY_CAPPED(float, kBlockDim);
42
using DI = Rebind<int, D>;
43
44
// Converts one vector's worth of random bits to floats in [1, 2).
45
// NOTE: as the convolution kernel sums to 0, it doesn't matter if inputs are in
46
// [0, 1) or in [1, 2).
47
void BitsToFloat(const uint32_t* JXL_RESTRICT random_bits,
48
144M
                 float* JXL_RESTRICT floats) {
49
144M
  const HWY_FULL(float) df;
50
144M
  const HWY_FULL(uint32_t) du;
51
52
144M
  const auto bits = Load(du, random_bits);
53
  // 1.0 + 23 random mantissa bits = [1, 2)
54
144M
  const auto rand12 = BitCast(df, Or(ShiftRight<9>(bits), Set(du, 0x3F800000)));
55
144M
  Store(rand12, df, floats);
56
144M
}
jxl::N_SSE4::BitsToFloat(unsigned int const*, float*)
Line
Count
Source
48
25.6M
                 float* JXL_RESTRICT floats) {
49
25.6M
  const HWY_FULL(float) df;
50
25.6M
  const HWY_FULL(uint32_t) du;
51
52
25.6M
  const auto bits = Load(du, random_bits);
53
  // 1.0 + 23 random mantissa bits = [1, 2)
54
25.6M
  const auto rand12 = BitCast(df, Or(ShiftRight<9>(bits), Set(du, 0x3F800000)));
55
25.6M
  Store(rand12, df, floats);
56
25.6M
}
jxl::N_AVX2::BitsToFloat(unsigned int const*, float*)
Line
Count
Source
48
82.2M
                 float* JXL_RESTRICT floats) {
49
82.2M
  const HWY_FULL(float) df;
50
82.2M
  const HWY_FULL(uint32_t) du;
51
52
82.2M
  const auto bits = Load(du, random_bits);
53
  // 1.0 + 23 random mantissa bits = [1, 2)
54
82.2M
  const auto rand12 = BitCast(df, Or(ShiftRight<9>(bits), Set(du, 0x3F800000)));
55
82.2M
  Store(rand12, df, floats);
56
82.2M
}
jxl::N_SSE2::BitsToFloat(unsigned int const*, float*)
Line
Count
Source
48
36.6M
                 float* JXL_RESTRICT floats) {
49
36.6M
  const HWY_FULL(float) df;
50
36.6M
  const HWY_FULL(uint32_t) du;
51
52
36.6M
  const auto bits = Load(du, random_bits);
53
  // 1.0 + 23 random mantissa bits = [1, 2)
54
36.6M
  const auto rand12 = BitCast(df, Or(ShiftRight<9>(bits), Set(du, 0x3F800000)));
55
36.6M
  Store(rand12, df, floats);
56
36.6M
}
57
58
void RandomImage(Xorshift128Plus* rng, const Rect& rect,
59
214k
                 ImageF* JXL_RESTRICT noise) {
60
214k
  const size_t xsize = rect.xsize();
61
214k
  const size_t ysize = rect.ysize();
62
63
  // May exceed the vector size, hence we have two loops over x below.
64
214k
  constexpr size_t kFloatsPerBatch =
65
214k
      Xorshift128Plus::N * sizeof(uint64_t) / sizeof(float);
66
214k
  HWY_ALIGN uint64_t batch64[Xorshift128Plus::N] = {};
67
214k
  HWY_ALIGN uint32_t batch32[2 * Xorshift128Plus::N];
68
69
214k
  const HWY_FULL(float) df;
70
214k
  const size_t N = Lanes(df);
71
72
7.66M
  for (size_t y = 0; y < ysize; ++y) {
73
7.44M
    float* JXL_RESTRICT row = rect.Row(noise, y);
74
75
7.44M
    size_t x = 0;
76
    // Only entire batches (avoids exceeding the image padding).
77
62.4M
    for (; x + kFloatsPerBatch < xsize; x += kFloatsPerBatch) {
78
54.9M
      rng->Fill(batch64);
79
      // Workaround for https://github.com/llvm/llvm-project/issues/121229
80
54.9M
      memcpy(batch32, batch64, sizeof(batch32));
81
197M
      for (size_t i = 0; i < kFloatsPerBatch; i += Lanes(df)) {
82
142M
        BitsToFloat(batch32 + i, row + x + i);
83
142M
      }
84
54.9M
    }
85
86
    // Any remaining pixels, rounded up to vectors (safe due to padding).
87
7.44M
    rng->Fill(batch64);
88
    // Workaround for https://github.com/llvm/llvm-project/issues/121229
89
7.44M
    memcpy(batch32, batch64, sizeof(batch32));
90
7.44M
    size_t batch_pos = 0;  // < kFloatsPerBatch
91
17.9M
    for (; x < xsize; x += N) {
92
10.4M
      BitsToFloat(batch32 + batch_pos, row + x);
93
10.4M
      batch_pos += N;
94
10.4M
    }
95
7.44M
  }
96
214k
}
dec_noise.cc:jxl::N_SSE4::RandomImage(jxl::N_SSE4::(anonymous namespace)::Xorshift128Plus*, jxl::RectT<unsigned long> const&, jxl::Plane<float>*)
Line
Count
Source
59
39.6k
                 ImageF* JXL_RESTRICT noise) {
60
39.6k
  const size_t xsize = rect.xsize();
61
39.6k
  const size_t ysize = rect.ysize();
62
63
  // May exceed the vector size, hence we have two loops over x below.
64
39.6k
  constexpr size_t kFloatsPerBatch =
65
39.6k
      Xorshift128Plus::N * sizeof(uint64_t) / sizeof(float);
66
39.6k
  HWY_ALIGN uint64_t batch64[Xorshift128Plus::N] = {};
67
39.6k
  HWY_ALIGN uint32_t batch32[2 * Xorshift128Plus::N];
68
69
39.6k
  const HWY_FULL(float) df;
70
39.6k
  const size_t N = Lanes(df);
71
72
822k
  for (size_t y = 0; y < ysize; ++y) {
73
782k
    float* JXL_RESTRICT row = rect.Row(noise, y);
74
75
782k
    size_t x = 0;
76
    // Only entire batches (avoids exceeding the image padding).
77
7.21M
    for (; x + kFloatsPerBatch < xsize; x += kFloatsPerBatch) {
78
6.43M
      rng->Fill(batch64);
79
      // Workaround for https://github.com/llvm/llvm-project/issues/121229
80
6.43M
      memcpy(batch32, batch64, sizeof(batch32));
81
32.1M
      for (size_t i = 0; i < kFloatsPerBatch; i += Lanes(df)) {
82
25.7M
        BitsToFloat(batch32 + i, row + x + i);
83
25.7M
      }
84
6.43M
    }
85
86
    // Any remaining pixels, rounded up to vectors (safe due to padding).
87
782k
    rng->Fill(batch64);
88
    // Workaround for https://github.com/llvm/llvm-project/issues/121229
89
782k
    memcpy(batch32, batch64, sizeof(batch32));
90
782k
    size_t batch_pos = 0;  // < kFloatsPerBatch
91
2.28M
    for (; x < xsize; x += N) {
92
1.49M
      BitsToFloat(batch32 + batch_pos, row + x);
93
1.49M
      batch_pos += N;
94
1.49M
    }
95
782k
  }
96
39.6k
}
dec_noise.cc:jxl::N_AVX2::RandomImage(jxl::N_AVX2::(anonymous namespace)::Xorshift128Plus*, jxl::RectT<unsigned long> const&, jxl::Plane<float>*)
Line
Count
Source
59
142k
                 ImageF* JXL_RESTRICT noise) {
60
142k
  const size_t xsize = rect.xsize();
61
142k
  const size_t ysize = rect.ysize();
62
63
  // May exceed the vector size, hence we have two loops over x below.
64
142k
  constexpr size_t kFloatsPerBatch =
65
142k
      Xorshift128Plus::N * sizeof(uint64_t) / sizeof(float);
66
142k
  HWY_ALIGN uint64_t batch64[Xorshift128Plus::N] = {};
67
142k
  HWY_ALIGN uint32_t batch32[2 * Xorshift128Plus::N];
68
69
142k
  const HWY_FULL(float) df;
70
142k
  const size_t N = Lanes(df);
71
72
5.21M
  for (size_t y = 0; y < ysize; ++y) {
73
5.07M
    float* JXL_RESTRICT row = rect.Row(noise, y);
74
75
5.07M
    size_t x = 0;
76
    // Only entire batches (avoids exceeding the image padding).
77
43.3M
    for (; x + kFloatsPerBatch < xsize; x += kFloatsPerBatch) {
78
38.2M
      rng->Fill(batch64);
79
      // Workaround for https://github.com/llvm/llvm-project/issues/121229
80
38.2M
      memcpy(batch32, batch64, sizeof(batch32));
81
114M
      for (size_t i = 0; i < kFloatsPerBatch; i += Lanes(df)) {
82
76.5M
        BitsToFloat(batch32 + i, row + x + i);
83
76.5M
      }
84
38.2M
    }
85
86
    // Any remaining pixels, rounded up to vectors (safe due to padding).
87
5.07M
    rng->Fill(batch64);
88
    // Workaround for https://github.com/llvm/llvm-project/issues/121229
89
5.07M
    memcpy(batch32, batch64, sizeof(batch32));
90
5.07M
    size_t batch_pos = 0;  // < kFloatsPerBatch
91
11.6M
    for (; x < xsize; x += N) {
92
6.58M
      BitsToFloat(batch32 + batch_pos, row + x);
93
6.58M
      batch_pos += N;
94
6.58M
    }
95
5.07M
  }
96
142k
}
dec_noise.cc:jxl::N_SSE2::RandomImage(jxl::N_SSE2::(anonymous namespace)::Xorshift128Plus*, jxl::RectT<unsigned long> const&, jxl::Plane<float>*)
Line
Count
Source
59
32.4k
                 ImageF* JXL_RESTRICT noise) {
60
32.4k
  const size_t xsize = rect.xsize();
61
32.4k
  const size_t ysize = rect.ysize();
62
63
  // May exceed the vector size, hence we have two loops over x below.
64
32.4k
  constexpr size_t kFloatsPerBatch =
65
32.4k
      Xorshift128Plus::N * sizeof(uint64_t) / sizeof(float);
66
32.4k
  HWY_ALIGN uint64_t batch64[Xorshift128Plus::N] = {};
67
32.4k
  HWY_ALIGN uint32_t batch32[2 * Xorshift128Plus::N];
68
69
32.4k
  const HWY_FULL(float) df;
70
32.4k
  const size_t N = Lanes(df);
71
72
1.62M
  for (size_t y = 0; y < ysize; ++y) {
73
1.59M
    float* JXL_RESTRICT row = rect.Row(noise, y);
74
75
1.59M
    size_t x = 0;
76
    // Only entire batches (avoids exceeding the image padding).
77
11.8M
    for (; x + kFloatsPerBatch < xsize; x += kFloatsPerBatch) {
78
10.2M
      rng->Fill(batch64);
79
      // Workaround for https://github.com/llvm/llvm-project/issues/121229
80
10.2M
      memcpy(batch32, batch64, sizeof(batch32));
81
50.9M
      for (size_t i = 0; i < kFloatsPerBatch; i += Lanes(df)) {
82
40.7M
        BitsToFloat(batch32 + i, row + x + i);
83
40.7M
      }
84
10.2M
    }
85
86
    // Any remaining pixels, rounded up to vectors (safe due to padding).
87
1.59M
    rng->Fill(batch64);
88
    // Workaround for https://github.com/llvm/llvm-project/issues/121229
89
1.59M
    memcpy(batch32, batch64, sizeof(batch32));
90
1.59M
    size_t batch_pos = 0;  // < kFloatsPerBatch
91
3.98M
    for (; x < xsize; x += N) {
92
2.39M
      BitsToFloat(batch32 + batch_pos, row + x);
93
2.39M
      batch_pos += N;
94
2.39M
    }
95
1.59M
  }
96
32.4k
}
97
void Random3Planes(size_t visible_frame_index, size_t nonvisible_frame_index,
98
                   size_t x0, size_t y0, const std::pair<ImageF*, Rect>& plane0,
99
                   const std::pair<ImageF*, Rect>& plane1,
100
71.4k
                   const std::pair<ImageF*, Rect>& plane2) {
101
71.4k
  HWY_ALIGN Xorshift128Plus rng(visible_frame_index, nonvisible_frame_index, x0,
102
71.4k
                                y0);
103
71.4k
  RandomImage(&rng, plane0.second, plane0.first);
104
71.4k
  RandomImage(&rng, plane1.second, plane1.first);
105
71.4k
  RandomImage(&rng, plane2.second, plane2.first);
106
71.4k
}
jxl::N_SSE4::Random3Planes(unsigned long, unsigned long, unsigned long, unsigned long, std::__1::pair<jxl::Plane<float>*, jxl::RectT<unsigned long> > const&, std::__1::pair<jxl::Plane<float>*, jxl::RectT<unsigned long> > const&, std::__1::pair<jxl::Plane<float>*, jxl::RectT<unsigned long> > const&)
Line
Count
Source
100
13.2k
                   const std::pair<ImageF*, Rect>& plane2) {
101
13.2k
  HWY_ALIGN Xorshift128Plus rng(visible_frame_index, nonvisible_frame_index, x0,
102
13.2k
                                y0);
103
13.2k
  RandomImage(&rng, plane0.second, plane0.first);
104
13.2k
  RandomImage(&rng, plane1.second, plane1.first);
105
13.2k
  RandomImage(&rng, plane2.second, plane2.first);
106
13.2k
}
jxl::N_AVX2::Random3Planes(unsigned long, unsigned long, unsigned long, unsigned long, std::__1::pair<jxl::Plane<float>*, jxl::RectT<unsigned long> > const&, std::__1::pair<jxl::Plane<float>*, jxl::RectT<unsigned long> > const&, std::__1::pair<jxl::Plane<float>*, jxl::RectT<unsigned long> > const&)
Line
Count
Source
100
47.3k
                   const std::pair<ImageF*, Rect>& plane2) {
101
47.3k
  HWY_ALIGN Xorshift128Plus rng(visible_frame_index, nonvisible_frame_index, x0,
102
47.3k
                                y0);
103
47.3k
  RandomImage(&rng, plane0.second, plane0.first);
104
47.3k
  RandomImage(&rng, plane1.second, plane1.first);
105
47.3k
  RandomImage(&rng, plane2.second, plane2.first);
106
47.3k
}
jxl::N_SSE2::Random3Planes(unsigned long, unsigned long, unsigned long, unsigned long, std::__1::pair<jxl::Plane<float>*, jxl::RectT<unsigned long> > const&, std::__1::pair<jxl::Plane<float>*, jxl::RectT<unsigned long> > const&, std::__1::pair<jxl::Plane<float>*, jxl::RectT<unsigned long> > const&)
Line
Count
Source
100
10.8k
                   const std::pair<ImageF*, Rect>& plane2) {
101
10.8k
  HWY_ALIGN Xorshift128Plus rng(visible_frame_index, nonvisible_frame_index, x0,
102
10.8k
                                y0);
103
10.8k
  RandomImage(&rng, plane0.second, plane0.first);
104
10.8k
  RandomImage(&rng, plane1.second, plane1.first);
105
10.8k
  RandomImage(&rng, plane2.second, plane2.first);
106
10.8k
}
107
108
// NOLINTNEXTLINE(google-readability-namespace-comments)
109
}  // namespace HWY_NAMESPACE
110
}  // namespace jxl
111
HWY_AFTER_NAMESPACE();
112
113
#if HWY_ONCE
114
namespace jxl {
115
116
namespace {
117
HWY_EXPORT(Random3Planes);
118
}  // namespace
119
120
void PrepareNoiseInput(const PassesDecoderState& dec_state,
121
                       const FrameDimensions& frame_dim,
122
                       const FrameHeader& frame_header, size_t group_index,
123
14.5k
                       size_t thread) {
124
14.5k
  size_t group_dim = frame_dim.group_dim;
125
14.5k
  const size_t gx = group_index % frame_dim.xsize_groups;
126
14.5k
  const size_t gy = group_index / frame_dim.xsize_groups;
127
14.5k
  RenderPipelineInput input =
128
14.5k
      dec_state.render_pipeline->GetInputBuffers(group_index, thread);
129
14.5k
  size_t noise_c_start =
130
14.5k
      3 + frame_header.nonserialized_metadata->m.num_extra_channels;
131
  // When the color channels are downsampled, we need to generate more noise
132
  // input for the current group than just the group dimensions.
133
14.5k
  std::pair<ImageF*, Rect> rects[3];
134
38.6k
  for (size_t iy = 0; iy < frame_header.upsampling; iy++) {
135
95.5k
    for (size_t ix = 0; ix < frame_header.upsampling; ix++) {
136
285k
      for (size_t c = 0; c < 3; c++) {
137
214k
        auto r = input.GetBuffer(noise_c_start + c);
138
214k
        rects[c].first = r.first;
139
214k
        size_t x1 = r.second.x0() + r.second.xsize();
140
214k
        size_t y1 = r.second.y0() + r.second.ysize();
141
214k
        rects[c].second =
142
214k
            Rect(r.second.x0() + ix * group_dim, r.second.y0() + iy * group_dim,
143
214k
                 group_dim, group_dim, x1, y1);
144
214k
      }
145
71.4k
      HWY_DYNAMIC_DISPATCH(Random3Planes)
146
71.4k
      (dec_state.visible_frame_index, dec_state.nonvisible_frame_index,
147
71.4k
       (gx * frame_header.upsampling + ix) * group_dim,
148
71.4k
       (gy * frame_header.upsampling + iy) * group_dim, rects[0], rects[1],
149
71.4k
       rects[2]);
150
71.4k
    }
151
24.1k
  }
152
14.5k
}
153
154
104k
void DecodeFloatParam(float precision, float* val, BitReader* br) {
155
104k
  const int absval_quant = br->ReadFixedBits<10>();
156
104k
  *val = absval_quant / precision;
157
104k
}
158
159
13.0k
Status DecodeNoise(BitReader* br, NoiseParams* noise_params) {
160
104k
  for (float& i : noise_params->lut) {
161
104k
    DecodeFloatParam(kNoisePrecision, &i, br);
162
104k
  }
163
13.0k
  return true;
164
13.0k
}
165
166
}  // namespace jxl
167
#endif  // HWY_ONCE