summaryrefslogtreecommitdiffstats
path: root/third_party/jpeg-xl/lib/jxl/render_pipeline/render_pipeline.h
blob: fa7a1e58c409300f1c1fa7c7878dfeffe820ab37 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
// 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_H_
#define LIB_JXL_RENDER_PIPELINE_RENDER_PIPELINE_H_

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <utility>
#include <vector>

#include "lib/jxl/base/rect.h"
#include "lib/jxl/base/status.h"
#include "lib/jxl/frame_dimensions.h"
#include "lib/jxl/image.h"
#include "lib/jxl/render_pipeline/render_pipeline_stage.h"

namespace jxl {

// Interface to provide input to the rendering pipeline. When this object is
// destroyed, all the data in the provided ImageF's Rects must have been
// initialized.
class RenderPipelineInput {
 public:
  RenderPipelineInput(const RenderPipelineInput&) = delete;
  RenderPipelineInput(RenderPipelineInput&& other) noexcept {
    *this = std::move(other);
  }
  RenderPipelineInput& operator=(RenderPipelineInput&& other) noexcept {
    pipeline_ = other.pipeline_;
    group_id_ = other.group_id_;
    thread_id_ = other.thread_id_;
    buffers_ = std::move(other.buffers_);
    other.pipeline_ = nullptr;
    return *this;
  }

  RenderPipelineInput() = default;
  Status Done();

  const std::pair<ImageF*, Rect>& GetBuffer(size_t c) const {
    JXL_ASSERT(c < buffers_.size());
    return buffers_[c];
  }

 private:
  RenderPipeline* pipeline_ = nullptr;
  size_t group_id_;
  size_t thread_id_;
  std::vector<std::pair<ImageF*, Rect>> buffers_;
  friend class RenderPipeline;
};

class RenderPipeline {
 public:
  class Builder {
   public:
    explicit Builder(size_t num_c) : num_c_(num_c) { JXL_ASSERT(num_c > 0); }

    // Adds a stage to the pipeline. Must be called at least once; the last
    // added stage cannot have kInOut channels.
    void AddStage(std::unique_ptr<RenderPipelineStage> stage);

    // Enables using the simple (i.e. non-memory-efficient) implementation of
    // the pipeline.
    void UseSimpleImplementation() { use_simple_implementation_ = true; }

    // Finalizes setup of the pipeline. Shifts for all channels should be 0 at
    // this point.
    StatusOr<std::unique_ptr<RenderPipeline>> Finalize(
        FrameDimensions frame_dimensions) &&;

   private:
    std::vector<std::unique_ptr<RenderPipelineStage>> stages_;
    size_t num_c_;
    bool use_simple_implementation_ = false;
  };

  friend class Builder;

  virtual ~RenderPipeline() = default;

  Status IsInitialized() const {
    for (const auto& stage : stages_) {
      JXL_RETURN_IF_ERROR(stage->IsInitialized());
    }
    return true;
  }

  // Allocates storage to run with `num` threads. If `use_group_ids` is true,
  // storage is allocated for each group, not each thread. The behaviour is
  // undefined if calling this function multiple times with a different value
  // for `use_group_ids`.
  Status PrepareForThreads(size_t num, bool use_group_ids);

  // Retrieves a buffer where input data should be stored by the callee. When
  // input has been provided for all buffers, the pipeline will complete its
  // processing. This method may be called multiple times concurrently from
  // different threads, provided that a different `thread_id` is given.
  RenderPipelineInput GetInputBuffers(size_t group_id, size_t thread_id);

  size_t PassesWithAllInput() const {
    return *std::min_element(group_completed_passes_.begin(),
                             group_completed_passes_.end());
  }

  virtual void ClearDone(size_t i) {}

 protected:
  std::vector<std::unique_ptr<RenderPipelineStage>> stages_;
  // Shifts for every channel at the input of each stage.
  std::vector<std::vector<std::pair<size_t, size_t>>> channel_shifts_;

  // Amount of (cumulative) padding required by each stage and channel, in
  // either direction.
  std::vector<std::vector<std::pair<size_t, size_t>>> padding_;

  FrameDimensions frame_dimensions_;

  std::vector<uint8_t> group_completed_passes_;

  friend class RenderPipelineInput;

 private:
  Status InputReady(size_t group_id, size_t thread_id,
                    const std::vector<std::pair<ImageF*, Rect>>& buffers);

  virtual std::vector<std::pair<ImageF*, Rect>> PrepareBuffers(
      size_t group_id, size_t thread_id) = 0;

  virtual Status ProcessBuffers(size_t group_id, size_t thread_id) = 0;

  // Note that this method may be called multiple times with different (or
  // equal) `num`.
  virtual Status PrepareForThreadsInternal(size_t num, bool use_group_ids) = 0;

  // Called once frame dimensions and stages are known.
  virtual Status Init() { return true; }
};

}  // namespace jxl

#endif  // LIB_JXL_RENDER_PIPELINE_RENDER_PIPELINE_H_