]> git.sesse.net Git - nageru/blob - quicksync_encoder_impl.h
Add functionality for MJPEG export.
[nageru] / quicksync_encoder_impl.h
1 #ifndef _QUICKSYNC_ENCODER_IMPL_H
2 #define _QUICKSYNC_ENCODER_IMPL_H 1
3
4 #include <epoxy/egl.h>
5 #include <movit/image_format.h>
6 #include <va/va.h>
7
8 #include <condition_variable>
9 #include <map>
10 #include <memory>
11 #include <mutex>
12 #include <queue>
13 #include <string>
14 #include <stack>
15 #include <thread>
16 #include <unordered_map>
17
18 #include "audio_encoder.h"
19 #include "defs.h"
20 #include "timebase.h"
21 #include "print_latency.h"
22 #include "ref_counted_gl_sync.h"
23 #include "va_display_with_cleanup.h"
24
25 #define SURFACE_NUM 16 /* 16 surfaces for source YUV */
26 #define MAX_NUM_REF1 16 // Seemingly a hardware-fixed value, not related to SURFACE_NUM
27 #define MAX_NUM_REF2 32 // Seemingly a hardware-fixed value, not related to SURFACE_NUM
28
29 struct __bitstream {
30     unsigned int *buffer;
31     int bit_offset;
32     int max_size_in_dword;
33 };
34 typedef struct __bitstream bitstream;
35
36 namespace movit {
37 class ResourcePool;
38 }
39 class DiskSpaceEstimator;
40 class QSurface;
41 class X264Encoder;
42
43 class QuickSyncEncoderImpl {
44 public:
45         QuickSyncEncoderImpl(const std::string &filename, movit::ResourcePool *resource_pool, QSurface *surface, const std::string &va_display, int width, int height, AVOutputFormat *oformat, X264Encoder *x264_encoder, DiskSpaceEstimator *disk_space_estimator);
46         ~QuickSyncEncoderImpl();
47         void add_audio(int64_t pts, std::vector<float> audio);
48         bool is_zerocopy() const;
49         bool begin_frame(int64_t pts, int64_t duration, movit::YCbCrLumaCoefficients ycbcr_coefficients, const std::vector<RefCountedFrame> &input_frames, GLuint *y_tex, GLuint *cbcr_tex);
50         RefCountedGLsync end_frame();
51         void shutdown();
52         void close_file();
53         void release_gl_resources();
54         void set_stream_mux(Mux *mux)
55         {
56                 stream_mux = mux;
57         }
58
59         // So we never get negative dts.
60         int64_t global_delay() const {
61                 return int64_t(ip_period - 1) * (TIMEBASE / MAX_FPS);
62         }
63
64 private:
65         struct storage_task {
66                 unsigned long long display_order;
67                 int frame_type;
68                 std::vector<float> audio;
69                 int64_t pts, dts, duration;
70                 movit::YCbCrLumaCoefficients ycbcr_coefficients;
71                 ReceivedTimestamps received_ts;
72                 std::vector<size_t> ref_display_frame_numbers;
73         };
74         struct PendingFrame {
75                 RefCountedGLsync fence;
76                 std::vector<RefCountedFrame> input_frames;
77                 int64_t pts, duration;
78                 movit::YCbCrLumaCoefficients ycbcr_coefficients;
79         };
80         struct GLSurface {
81                 // Only if x264_video_to_disk == false.
82                 VASurfaceID src_surface, ref_surface;
83                 VABufferID coded_buf;
84                 VAImage surface_image;
85
86                 // Only if use_zerocopy == true (which implies x264_video_to_disk == false).
87                 GLuint y_tex, cbcr_tex;
88                 EGLImage y_egl_image, cbcr_egl_image;
89
90                 // Only if use_zerocopy == false.
91                 GLuint pbo;
92                 uint8_t *y_ptr, *cbcr_ptr;
93                 size_t y_offset, cbcr_offset;
94
95                 // Surfaces can be busy (have refcount > 0) for a variety of
96                 // reasons: First of all because they belong to a frame that's
97                 // under encoding. But also reference frames take refcounts;
98                 // while a frame is being encoded, all its reference frames
99                 // also have increased refcounts so that they are not dropped.
100                 // Similarly, just being in <reference_frames> increases the
101                 // refcount. Until it is back to zero, the surface cannot be given
102                 // out for encoding another frame. Use release_gl_surface()
103                 // to reduce the refcount, which will free the surface if
104                 // the refcount reaches zero.
105                 //
106                 // Protected by storage_task_queue_mutex.
107                 int refcount = 0;
108         };
109
110         void open_output_file(const std::string &filename);
111         void encode_thread_func();
112         void encode_remaining_frames_as_p(int encoding_frame_num, int gop_start_display_frame_num, int64_t last_dts);
113         void add_packet_for_uncompressed_frame(int64_t pts, int64_t duration, const uint8_t *data);
114         void pass_frame(PendingFrame frame, int display_frame_num, int64_t pts, int64_t duration);
115         void encode_frame(PendingFrame frame, int encoding_frame_num, int display_frame_num, int gop_start_display_frame_num,
116                           int frame_type, int64_t pts, int64_t dts, int64_t duration, movit::YCbCrLumaCoefficients ycbcr_coefficients);
117         void storage_task_thread();
118         void storage_task_enqueue(storage_task task);
119         void save_codeddata(GLSurface *surf, storage_task task);
120         int render_packedsequence(movit::YCbCrLumaCoefficients ycbcr_coefficients);
121         int render_packedpicture();
122         void render_packedslice();
123         int render_sequence();
124         int render_picture(GLSurface *surf, int frame_type, int display_frame_num, int gop_start_display_frame_num);
125         void sps_rbsp(movit::YCbCrLumaCoefficients ycbcr_coefficients, bitstream *bs);
126         void pps_rbsp(bitstream *bs);
127         int build_packed_pic_buffer(unsigned char **header_buffer);
128         int render_slice(int encoding_frame_num, int display_frame_num, int gop_start_display_frame_num, int frame_type);
129         void slice_header(bitstream *bs);
130         int build_packed_seq_buffer(movit::YCbCrLumaCoefficients ycbcr_coefficients, unsigned char **header_buffer);
131         int build_packed_slice_buffer(unsigned char **header_buffer);
132         int init_va(const std::string &va_display);
133         void enable_zerocopy_if_possible();
134         int setup_encode();
135         void release_encode();
136         void update_ReferenceFrames(int current_display_frame, int frame_type);
137         void update_RefPicList_P(VAPictureH264 RefPicList0_P[MAX_NUM_REF2]);
138         void update_RefPicList_B(VAPictureH264 RefPicList0_B[MAX_NUM_REF2], VAPictureH264 RefPicList1_B[MAX_NUM_REF2]);
139         GLSurface *allocate_gl_surface();
140         void release_gl_surface(size_t display_frame_num);
141
142         bool is_shutdown = false;
143         bool has_released_gl_resources = false;
144         std::atomic<bool> use_zerocopy{false};
145
146         std::thread encode_thread, storage_thread;
147
148         std::mutex storage_task_queue_mutex;
149         std::condition_variable storage_task_queue_changed;
150         std::queue<storage_task> storage_task_queue;  // protected by storage_task_queue_mutex
151         bool storage_thread_should_quit = false;  // protected by storage_task_queue_mutex
152
153         std::mutex frame_queue_mutex;
154         std::condition_variable frame_queue_nonempty;
155         bool encode_thread_should_quit = false;  // under frame_queue_mutex
156
157         int current_storage_frame;
158
159         PendingFrame current_video_frame;  // Used only between begin_frame() and end_frame().
160         std::queue<PendingFrame> pending_video_frames;  // under frame_queue_mutex
161         movit::ResourcePool *resource_pool;
162         QSurface *surface;
163
164         // Frames that are done rendering and passed on to x264 (if enabled),
165         // but have not been encoded by Quick Sync yet, and thus also not freed.
166         // The key is the display frame number.
167         std::map<int, PendingFrame> reorder_buffer;
168         int quicksync_encoding_frame_num = 0;
169
170         std::mutex file_audio_encoder_mutex;
171         std::unique_ptr<AudioEncoder> file_audio_encoder;
172
173         X264Encoder *x264_encoder;  // nullptr if not using x264.
174
175         Mux* stream_mux = nullptr;  // To HTTP.
176         std::unique_ptr<Mux> file_mux;  // To local disk.
177
178         // Encoder parameters
179         std::unique_ptr<VADisplayWithCleanup> va_dpy;
180         VAProfile h264_profile = (VAProfile)~0;
181         VAConfigAttrib config_attrib[VAConfigAttribTypeMax];
182         int config_attrib_num = 0, enc_packed_header_idx;
183
184         GLSurface gl_surfaces[SURFACE_NUM];
185
186         // For all frames in encoding (refcount > 0), a pointer into gl_surfaces
187         // for the surface used for that frame. Protected by storage_task_queue_mutex.
188         // The key is display frame number.
189         std::unordered_map<size_t, GLSurface *> surface_for_frame;
190
191         VAConfigID config_id;
192         VAContextID context_id;
193         VAEncSequenceParameterBufferH264 seq_param;
194         VAEncPictureParameterBufferH264 pic_param;
195         VAEncSliceParameterBufferH264 slice_param;
196         VAPictureH264 CurrentCurrPic;
197
198         struct ReferenceFrame {
199                 VAPictureH264 pic;
200                 int display_number;  // To track reference counts.
201         };
202         std::deque<ReferenceFrame> reference_frames;
203
204         // Static quality settings.
205         static constexpr unsigned int frame_bitrate = 15000000 / 60;  // Doesn't really matter; only initial_qp does.
206         static constexpr unsigned int num_ref_frames = 2;
207         static constexpr int initial_qp = 15;
208         static constexpr int minimal_qp = 0;
209         static constexpr int intra_period = 30;
210         static constexpr int intra_idr_period = MAX_FPS;  // About a second; more at lower frame rates. Not ideal.
211
212         // Quality settings that are meant to be static, but might be overridden
213         // by the profile.
214         int constraint_set_flag = 0;
215         int h264_packedheader = 0; /* support pack header? */
216         int h264_maxref = (1<<16|1);
217         int h264_entropy_mode = 1; /* cabac */
218         int ip_period = 3;
219
220         unsigned int current_ref_frame_num = 0;  // Encoding frame order within this GOP, sans B-frames.
221
222         int frame_width;
223         int frame_height;
224         int frame_width_mbaligned;
225         int frame_height_mbaligned;
226
227         DiskSpaceEstimator *disk_space_estimator;
228 };
229
230 #endif  // !defined(_QUICKSYNC_ENCODER_IMPL_H)