lib/jxl/render_pipeline/simple_render_pipeline.cc

   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/render_pipeline/simple_render_pipeline.h"
   7
   8 #include <hwy/base.h>
   9
  10 #include "lib/jxl/image_ops.h"
  11 #include "lib/jxl/render_pipeline/render_pipeline_stage.h"
  12 #include "lib/jxl/sanitizers.h"
  13
  14 namespace jxl {
  15
  16 void SimpleRenderPipeline::PrepareForThreadsInternal(size_t num,
  17                                                      bool use_group_ids) {
  18   if (!channel_data_.empty()) {
  19     return;
  20   }
  21   auto ch_size = [](size_t frame_size, size_t shift) {
  22     return DivCeil(frame_size, 1 << shift) + kRenderPipelineXOffset * 2;
  23   };
  24   for (size_t c = 0; c < channel_shifts_[0].size(); c++) {
  25     channel_data_.push_back(ImageF(
  26         ch_size(frame_dimensions_.xsize_upsampled, channel_shifts_[0][c].first),
  27         ch_size(frame_dimensions_.ysize_upsampled,
  28                 channel_shifts_[0][c].second)));
  29     msan::PoisonImage(channel_data_.back());
  30   }
  31 }
  32
  33 Rect SimpleRenderPipeline::MakeChannelRect(size_t group_id, size_t channel) {
  34   size_t base_color_shift =
  35       CeilLog2Nonzero(frame_dimensions_.xsize_upsampled_padded /
  36                       frame_dimensions_.xsize_padded);
  37
  38   const size_t gx = group_id % frame_dimensions_.xsize_groups;
  39   const size_t gy = group_id / frame_dimensions_.xsize_groups;
  40   size_t xgroupdim = (frame_dimensions_.group_dim << base_color_shift) >>
  41                      channel_shifts_[0][channel].first;
  42   size_t ygroupdim = (frame_dimensions_.group_dim << base_color_shift) >>
  43                      channel_shifts_[0][channel].second;
  44   return Rect(
  45       kRenderPipelineXOffset + gx * xgroupdim,
  46       kRenderPipelineXOffset + gy * ygroupdim, xgroupdim, ygroupdim,
  47       kRenderPipelineXOffset + DivCeil(frame_dimensions_.xsize_upsampled,
  48                                        1 << channel_shifts_[0][channel].first),
  49       kRenderPipelineXOffset +
  50           DivCeil(frame_dimensions_.ysize_upsampled,
  51                   1 << channel_shifts_[0][channel].second));
  52 }
  53
  54 std::vector<std::pair<ImageF*, Rect>> SimpleRenderPipeline::PrepareBuffers(
  55     size_t group_id, size_t thread_id) {
  56   std::vector<std::pair<ImageF*, Rect>> ret;
  57   for (size_t c = 0; c < channel_data_.size(); c++) {
  58     ret.emplace_back(&channel_data_[c], MakeChannelRect(group_id, c));
  59   }
  60   return ret;
  61 }
  62
  63 void SimpleRenderPipeline::ProcessBuffers(size_t group_id, size_t thread_id) {
  64   for (size_t c = 0; c < channel_data_.size(); c++) {
  65     Rect r = MakeChannelRect(group_id, c);
  66     (void)r;
  67     JXL_CHECK_PLANE_INITIALIZED(channel_data_[c], r, c);
  68   }
  69
  70   if (PassesWithAllInput() <= processed_passes_) return;
  71   processed_passes_++;
  72
  73   for (size_t stage_id = 0; stage_id < stages_.size(); stage_id++) {
  74     const auto& stage = stages_[stage_id];
  75     // Prepare buffers for kInOut channels.
  76     std::vector<ImageF> new_channels(channel_data_.size());
  77     std::vector<ImageF*> output_channels(channel_data_.size());
  78
  79     std::vector<std::pair<size_t, size_t>> input_sizes(channel_data_.size());
  80     for (size_t c = 0; c < channel_data_.size(); c++) {
  81       input_sizes[c] =
  82           std::make_pair(channel_data_[c].xsize() - kRenderPipelineXOffset * 2,
  83                          channel_data_[c].ysize() - kRenderPipelineXOffset * 2);
  84     }
  85
  86     for (size_t c = 0; c < channel_data_.size(); c++) {
  87       if (stage->GetChannelMode(c) != RenderPipelineChannelMode::kInOut) {
  88         continue;
  89       }
  90       // Ensure that the newly allocated channels are large enough to avoid
  91       // problems with padding.
  92       new_channels[c] =
  93           ImageF(frame_dimensions_.xsize_upsampled_padded +
  94                      kRenderPipelineXOffset * 2 + hwy::kMaxVectorSize * 8,
  95                  frame_dimensions_.ysize_upsampled_padded +
  96                      kRenderPipelineXOffset * 2);
  97       new_channels[c].ShrinkTo(
  98           (input_sizes[c].first << stage->settings_.shift_x) +
  99               kRenderPipelineXOffset * 2,
 100           (input_sizes[c].second << stage->settings_.shift_y) +
 101               kRenderPipelineXOffset * 2);
 102       output_channels[c] = &new_channels[c];
 103     }
 104
 105     auto get_row = [&](size_t c, int64_t y) {
 106       return channel_data_[c].Row(kRenderPipelineXOffset + y) +
 107              kRenderPipelineXOffset;
 108     };
 109
 110     // Add mirrored pixes to all kInOut channels.
 111     for (size_t c = 0; c < channel_data_.size(); c++) {
 112       if (stage->GetChannelMode(c) != RenderPipelineChannelMode::kInOut) {
 113         continue;
 114       }
 115       // Horizontal mirroring.
 116       for (size_t y = 0; y < input_sizes[c].second; y++) {
 117         float* row = get_row(c, y);
 118         for (size_t ix = 0; ix < stage->settings_.border_x; ix++) {
 119           *(row - ix - 1) = row[Mirror(-ssize_t(ix) - 1, input_sizes[c].first)];
 120         }
 121         for (size_t ix = 0; ix < stage->settings_.border_x; ix++) {
 122           *(row + ix + input_sizes[c].first) =
 123               row[Mirror(ix + input_sizes[c].first, input_sizes[c].first)];
 124         }
 125       }
 126       // Vertical mirroring.
 127       for (int y = 0; y < static_cast<int>(stage->settings_.border_y); y++) {
 128         memcpy(get_row(c, -y - 1) - stage->settings_.border_x,
 129                get_row(c, Mirror(-ssize_t(y) - 1, input_sizes[c].second)) -
 130                    stage->settings_.border_x,
 131                sizeof(float) *
 132                    (input_sizes[c].first + 2 * stage->settings_.border_x));
 133       }
 134       for (int y = 0; y < static_cast<int>(stage->settings_.border_y); y++) {
 135         memcpy(
 136             get_row(c, input_sizes[c].second + y) - stage->settings_.border_x,
 137             get_row(c,
 138                     Mirror(input_sizes[c].second + y, input_sizes[c].second)) -
 139                 stage->settings_.border_x,
 140             sizeof(float) *
 141                 (input_sizes[c].first + 2 * stage->settings_.border_x));
 142       }
 143     }
 144
 145     size_t ysize = 0;
 146     size_t xsize = 0;
 147     for (size_t c = 0; c < channel_data_.size(); c++) {
 148       if (stage->GetChannelMode(c) == RenderPipelineChannelMode::kIgnored) {
 149         continue;
 150       }
 151       ysize = std::max(input_sizes[c].second, ysize);
 152       xsize = std::max(input_sizes[c].first, xsize);
 153     }
 154
 155     JXL_ASSERT(ysize != 0);
 156     JXL_ASSERT(xsize != 0);
 157
 158     RenderPipelineStage::RowInfo input_rows(channel_data_.size());
 159     RenderPipelineStage::RowInfo output_rows(channel_data_.size());
 160
 161     // Run the pipeline.
 162     {
 163       stage->SetInputSizes(input_sizes);
 164       int border_y = stage->settings_.border_y;
 165       for (size_t y = 0; y < ysize; y++) {
 166         // Prepare input rows.
 167         for (size_t c = 0; c < channel_data_.size(); c++) {
 168           if (stage->GetChannelMode(c) == RenderPipelineChannelMode::kIgnored) {
 169             continue;
 170           }
 171           input_rows[c].resize(2 * border_y + 1);
 172           for (int iy = -border_y; iy <= border_y; iy++) {
 173             input_rows[c][iy + border_y] =
 174                 channel_data_[c].Row(y + kRenderPipelineXOffset + iy);
 175           }
 176         }
 177         // Prepare output rows.
 178         for (size_t c = 0; c < channel_data_.size(); c++) {
 179           if (!output_channels[c]) continue;
 180           output_rows[c].resize(1 << stage->settings_.shift_y);
 181           for (size_t iy = 0; iy < output_rows[c].size(); iy++) {
 182             output_rows[c][iy] = output_channels[c]->Row(
 183                 (y << stage->settings_.shift_y) + iy + kRenderPipelineXOffset);
 184           }
 185         }
 186         stage->ProcessRow(input_rows, output_rows, /*xextra=*/0, xsize,
 187                           /*xpos=*/0, y, thread_id);
 188       }
 189     }
 190
 191     // Move new channels to current channels.
 192     for (size_t c = 0; c < channel_data_.size(); c++) {
 193       if (stage->GetChannelMode(c) != RenderPipelineChannelMode::kInOut) {
 194         continue;
 195       }
 196       channel_data_[c] = std::move(new_channels[c]);
 197     }
 198     for (size_t c = 0; c < channel_data_.size(); c++) {
 199       size_t next_stage = std::min(stage_id + 1, channel_shifts_.size() - 1);
 200       size_t xsize = DivCeil(frame_dimensions_.xsize_upsampled,
 201                              1 << channel_shifts_[next_stage][c].first);
 202       size_t ysize = DivCeil(frame_dimensions_.ysize_upsampled,
 203                              1 << channel_shifts_[next_stage][c].second);
 204       channel_data_[c].ShrinkTo(xsize + 2 * kRenderPipelineXOffset,
 205                                 ysize + 2 * kRenderPipelineXOffset);
 206       JXL_CHECK_PLANE_INITIALIZED(
 207           channel_data_[c],
 208           Rect(kRenderPipelineXOffset, kRenderPipelineXOffset, xsize, ysize),
 209           c);
 210     }
 211
 212     if (stage->SwitchToImageDimensions()) {
 213       size_t image_xsize, image_ysize;
 214       FrameOrigin frame_origin;
 215       stage->GetImageDimensions(&image_xsize, &image_ysize, &frame_origin);
 216       frame_dimensions_.Set(image_xsize, image_ysize, 0, 0, 0, false, 1);
 217       std::vector<ImageF> old_channels = std::move(channel_data_);
 218       channel_data_.clear();
 219       channel_data_.reserve(old_channels.size());
 220       for (size_t c = 0; c < old_channels.size(); c++) {
 221         channel_data_.emplace_back(2 * kRenderPipelineXOffset + image_xsize,
 222                                    2 * kRenderPipelineXOffset + image_ysize);
 223       }
 224       for (size_t y = 0; y < image_ysize; ++y) {
 225         for (size_t c = 0; c < channel_data_.size(); c++) {
 226           output_rows[c].resize(1);
 227           output_rows[c][0] = channel_data_[c].Row(kRenderPipelineXOffset + y);
 228         }
 229         // TODO(sboukortt): consider doing this only on the parts of the
 230         // background that won't be occluded.
 231         stage->ProcessPaddingRow(output_rows, image_xsize, 0, y);
 232       }
 233       ssize_t x0 = frame_origin.x0;
 234       ssize_t y0 = frame_origin.y0;
 235       size_t x0_fg = 0;
 236       size_t y0_fg = 0;
 237       if (x0 < 0) {
 238         xsize += x0;
 239         x0_fg -= x0;
 240         x0 = 0;
 241       }
 242       if (x0 + xsize > image_xsize) {
 243         xsize = image_xsize - x0;
 244       }
 245       if (y0 < 0) {
 246         ysize += y0;
 247         y0_fg -= x0;
 248         y0 = 0;
 249       }
 250       if (y0 + ysize > image_ysize) {
 251         ysize = image_ysize - y0;
 252       }
 253       const Rect rect_fg_relative_to_image =
 254           Rect(x0, y0, xsize, ysize)
 255               .Translate(kRenderPipelineXOffset, kRenderPipelineXOffset);
 256       const Rect rect_fg =
 257           Rect(x0_fg, y0_fg, xsize, ysize)
 258               .Translate(kRenderPipelineXOffset, kRenderPipelineXOffset);
 259       for (size_t c = 0; c < channel_data_.size(); c++) {
 260         CopyImageTo(rect_fg, old_channels[c], rect_fg_relative_to_image,
 261                     &channel_data_[c]);
 262       }
 263     }
 264   }
 265 }
 266 }  // namespace jxl