text: // Copyright (c) the JPEG XL Project Authors. All rights reserved.
  //
  // Use of this source code is governed by a BSD-style
  // license that can be found in the LICENSE file.
  
  #ifndef LIB_JXL_RENDER_PIPELINE_RENDER_PIPELINE_STAGE_H_
  #define LIB_JXL_RENDER_PIPELINE_RENDER_PIPELINE_STAGE_H_
  
  #include <stdint.h>
  
  #include "lib/jxl/base/arch_macros.h"
  #include "lib/jxl/frame_header.h"
  
  namespace jxl {
  
  // The first pixel in the input to RenderPipelineStage will be located at
  // this position. Pixels before this position may be accessed as padding.
  // This should be at least the RoundUpTo(maximum padding / 2, maximum vector
  // size) times 2: this is realized when using Gaborish + EPF + upsampling +
  // chroma subsampling.
  #if JXL_ARCH_ARM
  constexpr size_t kRenderPipelineXOffset = 16;
  #else
  constexpr size_t kRenderPipelineXOffset = 32;
  #endif
  
  enum class RenderPipelineChannelMode {
    // This channel is not modified by this stage.
    kIgnored = 0,
    // This channel is modified in-place.
    kInPlace = 1,
    // This channel is modified and written to a new buffer.
    kInOut = 2,
    // This channel is only read. These are the only stages that are assumed to
    // have observable effects, i.e. calls to ProcessRow for other stages may be
    // omitted if it can be shown they can't affect any kInput stage ProcessRow
    // call that happens inside image boundaries.
    kInput = 3,
  };
  
  class RenderPipeline;
  
  class RenderPipelineStage {
   protected:
    using Row = float*;
    using ChannelRows = std::vector<Row>;
  
   public:
    using RowInfo = std::vector<ChannelRows>;
    struct Settings {
      // Amount of padding required in the various directions by all channels
      // that have kInOut mode.
      size_t border_x = 0;
      size_t border_y = 0;
  
      // Log2 of the number of columns/rows of output that this stage will produce
      // for every input row for kInOut channels.
      size_t shift_x = 0;
      size_t shift_y = 0;
  
      static Settings ShiftX(size_t shift, size_t border) {
        Settings settings;
        settings.border_x = border;
        settings.shift_x = shift;
        return settings;
      }
  
      static Settings ShiftY(size_t shift, size_t border) {
        Settings settings;
        settings.border_y = border;
        settings.shift_y = shift;
        return settings;
      }
  
      static Settings Symmetric(size_t shift, size_t border) {
        Settings settings;
        settings.border_x = settings.border_y = border;
        settings.shift_x = settings.shift_y = shift;
        return settings;
      }
  
      static Settings SymmetricBorderOnly(size_t border) {
        return Symmetric(0, border);
      }
    };
  
    virtual ~RenderPipelineStage() = default;
  
    // Processes one row of input, producing the appropriate number of rows of
    // output. Input/output rows can be obtained by calls to
    // `GetInputRow`/`GetOutputRow`. `xsize+2*xextra` represents the total number
    // of pixels to be processed in the input row, where the first pixel is at
    // position `kRenderPipelineXOffset-xextra`. All pixels in the
    // `[kRenderPipelineXOffset-xextra-border_x,
    // kRenderPipelineXOffset+xsize+xextra+border_x)` range are initialized and
    // accessible. `xpos` and `ypos` represent the position of the first
    // (non-extra, i.e. in position kRenderPipelineXOffset) pixel in the center
    // row of the input in the full image. `xpos` is a multiple of
    // `GroupBorderAssigner::kPaddingXRound`. If `settings_.temp_buffer_size` is
    // nonzero, `temp` will point to an HWY-aligned buffer of at least that number
    // of floats; concurrent calls will have different buffers.
    virtual void ProcessRow(const RowInfo& input_rows, const RowInfo& output_rows,
                            size_t xextra, size_t xsize, size_t xpos, size_t ypos,
                            size_t thread_id) const = 0;
  
    // How each channel will be processed. Channels are numbered starting from
    // color channels (always 3) and followed by all other channels.
    virtual RenderPipelineChannelMode GetChannelMode(size_t c) const = 0;
  
   protected:
    explicit RenderPipelineStage(Settings settings) : settings_(settings) {}
  
    virtual Status IsInitialized() const { return true; }
  
    // Informs the stage about the total size of each channel. Few stages will
    // actually need to use this information.
    virtual void SetInputSizes(
        const std::vector<std::pair<size_t, size_t>>& input_sizes) {}
  
    virtual Status PrepareForThreads(size_t num_threads) { return true; }
  
    // Returns a pointer to the input row of channel `c` with offset `y`.
    // `y` must be in [-settings_.border_y, settings_.border_y]. `c` must be such
    // that `GetChannelMode(c) != kIgnored`. The returned pointer points to the
    // offset-ed row (i.e. kRenderPipelineXOffset has been applied).
    float* GetInputRow(const RowInfo& input_rows, size_t c, int offset) const {
      JXL_DASSERT(GetChannelMode(c) != RenderPipelineChannelMode::kIgnored);
      JXL_DASSERT(-offset <= static_cast<int>(settings_.border_y));
      JXL_DASSERT(offset <= static_cast<int>(settings_.border_y));
      return input_rows[c][settings_.border_y + offset] + kRenderPipelineXOffset;
    }
    // Similar to `GetInputRow`, but can only be used if `GetChannelMode(c) ==
    // kInOut`. Offset must be less than `1<<settings_.shift_y`.. The returned
    // pointer points to the offset-ed row (i.e. kRenderPipelineXOffset has been
    // applied).
    float* GetOutputRow(const RowInfo& output_rows, size_t c,
                        size_t offset) const {
      JXL_DASSERT(GetChannelMode(c) == RenderPipelineChannelMode::kInOut);
      JXL_DASSERT(offset <= 1ul << settings_.shift_y);
      return output_rows[c][offset] + kRenderPipelineXOffset;
    }
  
    // Indicates whether, from this stage on, the pipeline will operate on an
    // image- rather than frame-sized buffer. Only one stage in the pipeline
    // should return true, and it should implement ProcessPaddingRow below too.
    // It is assumed that, if there is a SwitchToImageDimensions() == true stage,
    // all kInput stages appear after it.
    virtual bool SwitchToImageDimensions() const { return false; }
  
    // If SwitchToImageDimensions returns true, then this should set xsize and
    // ysize to the image size, and frame_origin to the location of the frame
    // within the image. Otherwise, this is not called at all.
    virtual void GetImageDimensions(size_t* xsize, size_t* ysize,
                                    FrameOrigin* frame_origin) const {}
  
    // Produces the appropriate output data outside of the frame dimensions. xpos
    // and ypos are now relative to the full image.
    virtual void ProcessPaddingRow(const RowInfo& output_rows, size_t xsize,
                                   size_t xpos, size_t ypos) const {}
  
    virtual const char* GetName() const = 0;
  
    Settings settings_;
    friend class RenderPipeline;
    friend class SimpleRenderPipeline;
    friend class LowMemoryRenderPipeline;
  };
  
  }  // namespace jxl
  
  #endif  // LIB_JXL_RENDER_PIPELINE_RENDER_PIPELINE_STAGE_H_
  