test/encode/mpeg2enc.c

   1 /*
   2  * Copyright (c) 2012 Intel Corporation. All Rights Reserved.
   3  *
   4  * Permission is hereby granted, free of charge, to any person obtaining a
   5  * copy of this software and associated documentation files (the
   6  * "Software"), to deal in the Software without restriction, including
   7  * without limitation the rights to use, copy, modify, merge, publish,
   8  * distribute, sub license, and/or sell copies of the Software, and to
   9  * permit persons to whom the Software is furnished to do so, subject to
  10  * the following conditions:
  11  *
  12  * The above copyright notice and this permission notice (including the
  13  * next paragraph) shall be included in all copies or substantial portions
  14  * of the Software.
  15  *
  16  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
  17  * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  18  * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT.
  19  * IN NO EVENT SHALL PRECISION INSIGHT AND/OR ITS SUPPLIERS BE LIABLE FOR
  20  * ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
  21  * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
  22  * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  23  */
  24 /*
  25  * Simple MPEG-2 encoder based on libVA.
  26  *
  27  */
  28
  29 #include "sysdeps.h"
  30
  31 #include <getopt.h>
  32 #include <unistd.h>
  33
  34 #include <sys/time.h>
  35 #include <sys/types.h>
  36 #include <fcntl.h>
  37 #include <time.h>
  38 #include <pthread.h>
  39
  40 #include <va/va.h>
  41 #include <va/va_enc_mpeg2.h>
  42
  43 #include "va_display.h"
  44
  45 #define START_CODE_PICUTRE      0x00000100
  46 #define START_CODE_SLICE        0x00000101
  47 #define START_CODE_USER         0x000001B2
  48 #define START_CODE_SEQ          0x000001B3
  49 #define START_CODE_EXT          0x000001B5
  50 #define START_CODE_GOP          0x000001B8
  51
  52 #define CHROMA_FORMAT_RESERVED  0
  53 #define CHROMA_FORMAT_420       1
  54 #define CHROMA_FORMAT_422       2
  55 #define CHROMA_FORMAT_444       3
  56
  57 #define MAX_SLICES              128
  58
  59 enum {
  60     MPEG2_MODE_I = 0,
  61     MPEG2_MODE_IP,
  62     MPEG2_MODE_IPB,
  63 };
  64
  65 enum {
  66     MPEG2_LEVEL_LOW = 0,
  67     MPEG2_LEVEL_MAIN,
  68     MPEG2_LEVEL_HIGH,
  69 };
  70
  71 #define CHECK_VASTATUS(va_status, func)                                 \
  72     if (va_status != VA_STATUS_SUCCESS) {                               \
  73         fprintf(stderr, "%s:%s (%d) failed, exit\n", __func__, func, __LINE__); \
  74         exit(1);                                                        \
  75     }
  76
  77 static VAProfile mpeg2_va_profiles[] = {
  78     VAProfileMPEG2Simple,
  79     VAProfileMPEG2Main
  80 };
  81
  82 static struct _mpeg2_sampling_density
  83 {
  84     int samplers_per_line;
  85     int line_per_frame;
  86     int frame_per_sec;
  87 } mpeg2_upper_samplings[2][3] = {
  88     { { 0, 0, 0 },
  89       { 720, 576, 30 },
  90       { 0, 0, 0 },
  91     },
  92
  93     { { 352, 288, 30 },
  94       { 720, 576, 30 },
  95       { 1920, 1152, 60 },
  96     }
  97 };
  98
  99 struct mpeg2enc_context {
 100     /* args */
 101     int rate_control_mode;
 102     int fps;
 103     int mode; /* 0:I, 1:I/P, 2:I/P/B */
 104     VAProfile profile;
 105     int level;
 106     int width;
 107     int height;
 108     int frame_size;
 109     int num_pictures;
 110     int qp;
 111     FILE *ifp;
 112     FILE *ofp;
 113     unsigned char *frame_data_buffer;
 114     int intra_period;
 115     int ip_period;
 116     int bit_rate; /* in kbps */
 117     VAEncPictureType next_type;
 118     int next_display_order;
 119     int next_bframes;
 120     int new_sequence;
 121     int new_gop_header;
 122     int gop_header_in_display_order;
 123
 124     /* VA resource */
 125     VADisplay va_dpy;
 126     VAEncSequenceParameterBufferMPEG2 seq_param;
 127     VAEncPictureParameterBufferMPEG2 pic_param;
 128     VAEncSliceParameterBufferMPEG2 slice_param[MAX_SLICES];
 129     VAContextID context_id;
 130     VAConfigID config_id;
 131     VABufferID seq_param_buf_id;                /* Sequence level parameter */
 132     VABufferID pic_param_buf_id;                /* Picture level parameter */
 133     VABufferID slice_param_buf_id[MAX_SLICES];  /* Slice level parameter, multil slices */
 134     VABufferID codedbuf_buf_id;                 /* Output buffer, compressed data */
 135     VABufferID packed_seq_header_param_buf_id;
 136     VABufferID packed_seq_buf_id;
 137     VABufferID packed_pic_header_param_buf_id;
 138     VABufferID packed_pic_buf_id;
 139     int num_slice_groups;
 140     int codedbuf_i_size;
 141     int codedbuf_pb_size;
 142
 143     /* thread */
 144     pthread_t upload_thread_id;
 145     int upload_thread_value;
 146     int current_input_surface;
 147     int current_upload_surface;
 148 };
 149
 150 /*
 151  * mpeg2enc helpers
 152  */
 153 #define BITSTREAM_ALLOCATE_STEPPING     4096
 154
 155 struct __bitstream {
 156     unsigned int *buffer;
 157     int bit_offset;
 158     int max_size_in_dword;
 159 };
 160
 161 typedef struct __bitstream bitstream;
 162
 163 static unsigned int
 164 swap32(unsigned int val)
 165 {
 166     unsigned char *pval = (unsigned char *)&val;
 167
 168     return ((pval[0] << 24)     |
 169             (pval[1] << 16)     |
 170             (pval[2] << 8)      |
 171             (pval[3] << 0));
 172 }
 173
 174 static void
 175 bitstream_start(bitstream *bs)
 176 {
 177     bs->max_size_in_dword = BITSTREAM_ALLOCATE_STEPPING;
 178     bs->buffer = calloc(bs->max_size_in_dword * sizeof(int), 1);
 179     bs->bit_offset = 0;
 180 }
 181
 182 static void
 183 bitstream_end(bitstream *bs)
 184 {
 185     int pos = (bs->bit_offset >> 5);
 186     int bit_offset = (bs->bit_offset & 0x1f);
 187     int bit_left = 32 - bit_offset;
 188
 189     if (bit_offset) {
 190         bs->buffer[pos] = swap32((bs->buffer[pos] << bit_left));
 191     }
 192 }
 193
 194 static void
 195 bitstream_put_ui(bitstream *bs, unsigned int val, int size_in_bits)
 196 {
 197     int pos = (bs->bit_offset >> 5);
 198     int bit_offset = (bs->bit_offset & 0x1f);
 199     int bit_left = 32 - bit_offset;
 200
 201     if (!size_in_bits)
 202         return;
 203
 204     if (size_in_bits < 32)
 205         val &= ((1 << size_in_bits) - 1);
 206
 207     bs->bit_offset += size_in_bits;
 208
 209     if (bit_left > size_in_bits) {
 210         bs->buffer[pos] = (bs->buffer[pos] << size_in_bits | val);
 211     } else {
 212         size_in_bits -= bit_left;
 213         bs->buffer[pos] = (bs->buffer[pos] << bit_left) | (val >> size_in_bits);
 214         bs->buffer[pos] = swap32(bs->buffer[pos]);
 215
 216         if (pos + 1 == bs->max_size_in_dword) {
 217             bs->max_size_in_dword += BITSTREAM_ALLOCATE_STEPPING;
 218             bs->buffer = realloc(bs->buffer, bs->max_size_in_dword * sizeof(unsigned int));
 219         }
 220
 221         bs->buffer[pos + 1] = val;
 222     }
 223 }
 224
 225 static void
 226 bitstream_byte_aligning(bitstream *bs, int bit)
 227 {
 228     int bit_offset = (bs->bit_offset & 0x7);
 229     int bit_left = 8 - bit_offset;
 230     int new_val;
 231
 232     if (!bit_offset)
 233         return;
 234
 235     assert(bit == 0 || bit == 1);
 236
 237     if (bit)
 238         new_val = (1 << bit_left) - 1;
 239     else
 240         new_val = 0;
 241
 242     bitstream_put_ui(bs, new_val, bit_left);
 243 }
 244
 245 static struct mpeg2_frame_rate {
 246     int code;
 247     float value;
 248 } frame_rate_tab[] = {
 249     {1, 23.976},
 250     {2, 24.0},
 251     {3, 25.0},
 252     {4, 29.97},
 253     {5, 30},
 254     {6, 50},
 255     {7, 59.94},
 256     {8, 60}
 257 };
 258
 259 static int
 260 find_frame_rate_code(const VAEncSequenceParameterBufferMPEG2 *seq_param)
 261 {
 262     unsigned int delta = -1;
 263     int code = 1, i;
 264     float frame_rate_value = seq_param->frame_rate *
 265         (seq_param->sequence_extension.bits.frame_rate_extension_d + 1) /
 266         (seq_param->sequence_extension.bits.frame_rate_extension_n + 1);
 267
 268     for (i = 0; i < sizeof(frame_rate_tab) / sizeof(frame_rate_tab[0]); i++) {
 269
 270         if (abs(1000 * frame_rate_tab[i].value - 1000 * frame_rate_value) < delta) {
 271             code = frame_rate_tab[i].code;
 272             delta = abs(1000 * frame_rate_tab[i].value - 1000 * frame_rate_value);
 273         }
 274     }
 275
 276     return code;
 277 }
 278
 279 static void
 280 sps_rbsp(struct mpeg2enc_context *ctx,
 281          const VAEncSequenceParameterBufferMPEG2 *seq_param,
 282          bitstream *bs)
 283 {
 284     int frame_rate_code = find_frame_rate_code(seq_param);
 285
 286     if (ctx->new_sequence) {
 287         bitstream_put_ui(bs, START_CODE_SEQ, 32);
 288         bitstream_put_ui(bs, seq_param->picture_width, 12);
 289         bitstream_put_ui(bs, seq_param->picture_height, 12);
 290         bitstream_put_ui(bs, seq_param->aspect_ratio_information, 4);
 291         bitstream_put_ui(bs, frame_rate_code, 4); /* frame_rate_code */
 292         bitstream_put_ui(bs, (seq_param->bits_per_second + 399) / 400, 18); /* the low 18 bits of bit_rate */
 293         bitstream_put_ui(bs, 1, 1); /* marker_bit */
 294         bitstream_put_ui(bs, seq_param->vbv_buffer_size, 10);
 295         bitstream_put_ui(bs, 0, 1); /* constraint_parameter_flag, always 0 for MPEG-2 */
 296         bitstream_put_ui(bs, 0, 1); /* load_intra_quantiser_matrix */
 297         bitstream_put_ui(bs, 0, 1); /* load_non_intra_quantiser_matrix */
 298
 299         bitstream_byte_aligning(bs, 0);
 300
 301         bitstream_put_ui(bs, START_CODE_EXT, 32);
 302         bitstream_put_ui(bs, 1, 4); /* sequence_extension id */
 303         bitstream_put_ui(bs, seq_param->sequence_extension.bits.profile_and_level_indication, 8);
 304         bitstream_put_ui(bs, seq_param->sequence_extension.bits.progressive_sequence, 1);
 305         bitstream_put_ui(bs, seq_param->sequence_extension.bits.chroma_format, 2);
 306         bitstream_put_ui(bs, seq_param->picture_width >> 12, 2);
 307         bitstream_put_ui(bs, seq_param->picture_height >> 12, 2);
 308         bitstream_put_ui(bs, ((seq_param->bits_per_second + 399) / 400) >> 18, 12); /* bit_rate_extension */
 309         bitstream_put_ui(bs, 1, 1); /* marker_bit */
 310         bitstream_put_ui(bs, seq_param->vbv_buffer_size >> 10, 8);
 311         bitstream_put_ui(bs, seq_param->sequence_extension.bits.low_delay, 1);
 312         bitstream_put_ui(bs, seq_param->sequence_extension.bits.frame_rate_extension_n, 2);
 313         bitstream_put_ui(bs, seq_param->sequence_extension.bits.frame_rate_extension_d, 5);
 314
 315         bitstream_byte_aligning(bs, 0);
 316     }
 317
 318     if (ctx->new_gop_header) {
 319         bitstream_put_ui(bs, START_CODE_GOP, 32);
 320         bitstream_put_ui(bs, seq_param->gop_header.bits.time_code, 25);
 321         bitstream_put_ui(bs, seq_param->gop_header.bits.closed_gop, 1);
 322         bitstream_put_ui(bs, seq_param->gop_header.bits.broken_link, 1);
 323
 324         bitstream_byte_aligning(bs, 0);
 325     }
 326 }
 327
 328 static void
 329 pps_rbsp(const VAEncSequenceParameterBufferMPEG2 *seq_param,
 330          const VAEncPictureParameterBufferMPEG2 *pic_param,
 331          bitstream *bs)
 332 {
 333     int i;
 334     int chroma_420_type;
 335
 336     if (seq_param->sequence_extension.bits.chroma_format == CHROMA_FORMAT_420)
 337         chroma_420_type = pic_param->picture_coding_extension.bits.progressive_frame;
 338     else
 339         chroma_420_type = 0;
 340
 341     bitstream_put_ui(bs, START_CODE_PICUTRE, 32);
 342     bitstream_put_ui(bs, pic_param->temporal_reference, 10);
 343     bitstream_put_ui(bs,
 344                      pic_param->picture_type == VAEncPictureTypeIntra ? 1 :
 345                      pic_param->picture_type == VAEncPictureTypePredictive ? 2 : 3,
 346                      3);
 347     bitstream_put_ui(bs, 0xFFFF, 16); /* vbv_delay, always 0xFFFF */
 348
 349     if (pic_param->picture_type == VAEncPictureTypePredictive ||
 350         pic_param->picture_type == VAEncPictureTypeBidirectional) {
 351         bitstream_put_ui(bs, 0, 1); /* full_pel_forward_vector, always 0 for MPEG-2 */
 352         bitstream_put_ui(bs, 7, 3); /* forward_f_code, always 7 for MPEG-2 */
 353     }
 354
 355     if (pic_param->picture_type == VAEncPictureTypeBidirectional) {
 356         bitstream_put_ui(bs, 0, 1); /* full_pel_backward_vector, always 0 for MPEG-2 */
 357         bitstream_put_ui(bs, 7, 3); /* backward_f_code, always 7 for MPEG-2 */
 358     }
 359
 360     bitstream_put_ui(bs, 0, 1); /* extra_bit_picture, 0 */
 361
 362     bitstream_byte_aligning(bs, 0);
 363
 364     bitstream_put_ui(bs, START_CODE_EXT, 32);
 365     bitstream_put_ui(bs, 8, 4); /* Picture Coding Extension ID: 8 */
 366     bitstream_put_ui(bs, pic_param->f_code[0][0], 4);
 367     bitstream_put_ui(bs, pic_param->f_code[0][1], 4);
 368     bitstream_put_ui(bs, pic_param->f_code[1][0], 4);
 369     bitstream_put_ui(bs, pic_param->f_code[1][1], 4);
 370
 371     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.intra_dc_precision, 2);
 372     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.picture_structure, 2);
 373     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.top_field_first, 1);
 374     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.frame_pred_frame_dct, 1);
 375     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.concealment_motion_vectors, 1);
 376     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.q_scale_type, 1);
 377     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.intra_vlc_format, 1);
 378     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.alternate_scan, 1);
 379     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.repeat_first_field, 1);
 380     bitstream_put_ui(bs, chroma_420_type, 1);
 381     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.progressive_frame, 1);
 382     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.composite_display_flag, 1);
 383
 384     bitstream_byte_aligning(bs, 0);
 385
 386     if (pic_param->user_data_length) {
 387         bitstream_put_ui(bs, START_CODE_USER, 32);
 388
 389         for (i = 0; i < pic_param->user_data_length; i++) {
 390             bitstream_put_ui(bs, pic_param->user_data[i], 8);
 391         }
 392
 393         bitstream_byte_aligning(bs, 0);
 394     }
 395 }
 396
 397 static int
 398 build_packed_pic_buffer(const VAEncSequenceParameterBufferMPEG2 *seq_param,
 399                         const VAEncPictureParameterBufferMPEG2 *pic_param,
 400                         unsigned char **header_buffer)
 401 {
 402     bitstream bs;
 403
 404     bitstream_start(&bs);
 405     pps_rbsp(seq_param, pic_param, &bs);
 406     bitstream_end(&bs);
 407
 408     *header_buffer = (unsigned char *)bs.buffer;
 409     return bs.bit_offset;
 410 }
 411
 412 static int
 413 build_packed_seq_buffer(struct mpeg2enc_context *ctx,
 414                         const VAEncSequenceParameterBufferMPEG2 *seq_param,
 415                         unsigned char **header_buffer)
 416 {
 417     bitstream bs;
 418
 419     bitstream_start(&bs);
 420     sps_rbsp(ctx, seq_param, &bs);
 421     bitstream_end(&bs);
 422
 423     *header_buffer = (unsigned char *)bs.buffer;
 424     return bs.bit_offset;
 425 }
 426
 427 /*
 428  * mpeg2enc
 429  */
 430 #define SID_INPUT_PICTURE_0                     0
 431 #define SID_INPUT_PICTURE_1                     1
 432 #define SID_REFERENCE_PICTURE_L0                2
 433 #define SID_REFERENCE_PICTURE_L1                3
 434 #define SID_RECON_PICTURE                       4
 435 #define SID_NUMBER                              SID_RECON_PICTURE + 1
 436
 437 static VASurfaceID surface_ids[SID_NUMBER];
 438
 439 /*
 440  * upload thread function
 441  */
 442 static void *
 443 upload_yuv_to_surface(void *data)
 444 {
 445     struct mpeg2enc_context *ctx = data;
 446     VAImage surface_image;
 447     VAStatus va_status;
 448     void *surface_p = NULL;
 449     unsigned char *y_src, *u_src, *v_src;
 450     unsigned char *y_dst, *u_dst, *v_dst;
 451     int y_size = ctx->width * ctx->height;
 452     int u_size = (ctx->width >> 1) * (ctx->height >> 1);
 453     int row, col;
 454     size_t n_items;
 455
 456     do {
 457         n_items = fread(ctx->frame_data_buffer, ctx->frame_size, 1, ctx->ifp);
 458     } while (n_items != 1);
 459
 460     va_status = vaDeriveImage(ctx->va_dpy, surface_ids[ctx->current_upload_surface], &surface_image);
 461     CHECK_VASTATUS(va_status,"vaDeriveImage");
 462
 463     vaMapBuffer(ctx->va_dpy, surface_image.buf, &surface_p);
 464     assert(VA_STATUS_SUCCESS == va_status);
 465
 466     y_src = ctx->frame_data_buffer;
 467     u_src = ctx->frame_data_buffer + y_size; /* UV offset for NV12 */
 468     v_src = ctx->frame_data_buffer + y_size + u_size;
 469
 470     y_dst = surface_p + surface_image.offsets[0];
 471     u_dst = surface_p + surface_image.offsets[1]; /* UV offset for NV12 */
 472     v_dst = surface_p + surface_image.offsets[2];
 473
 474     /* Y plane */
 475     for (row = 0; row < surface_image.height; row++) {
 476         memcpy(y_dst, y_src, surface_image.width);
 477         y_dst += surface_image.pitches[0];
 478         y_src += ctx->width;
 479     }
 480
 481     if (surface_image.format.fourcc == VA_FOURCC_NV12) { /* UV plane */
 482         for (row = 0; row < surface_image.height / 2; row++) {
 483             for (col = 0; col < surface_image.width / 2; col++) {
 484                 u_dst[col * 2] = u_src[col];
 485                 u_dst[col * 2 + 1] = v_src[col];
 486             }
 487
 488             u_dst += surface_image.pitches[1];
 489             u_src += (ctx->width / 2);
 490             v_src += (ctx->width / 2);
 491         }
 492     } else {
 493         for (row = 0; row < surface_image.height / 2; row++) {
 494             for (col = 0; col < surface_image.width / 2; col++) {
 495                 u_dst[col] = u_src[col];
 496                 v_dst[col] = v_src[col];
 497             }
 498
 499             u_dst += surface_image.pitches[1];
 500             v_dst += surface_image.pitches[2];
 501             u_src += (ctx->width / 2);
 502             v_src += (ctx->width / 2);
 503         }
 504     }
 505
 506     vaUnmapBuffer(ctx->va_dpy, surface_image.buf);
 507     vaDestroyImage(ctx->va_dpy, surface_image.image_id);
 508
 509     return NULL;
 510 }
 511
 512 static void
 513 mpeg2enc_exit(struct mpeg2enc_context *ctx, int exit_code)
 514 {
 515     if (ctx->frame_data_buffer) {
 516         free(ctx->frame_data_buffer);
 517         ctx->frame_data_buffer = NULL;
 518     }
 519
 520     if (ctx->ifp) {
 521         fclose(ctx->ifp);
 522         ctx->ifp = NULL;
 523     }
 524
 525     if (ctx->ofp) {
 526         fclose(ctx->ofp);
 527         ctx->ofp = NULL;
 528     }
 529
 530     exit(exit_code);
 531 }
 532
 533 static void
 534 usage(char *program)
 535 {
 536     fprintf(stderr, "Usage: %s --help\n", program);
 537     fprintf(stderr, "\t--help   print this message\n");
 538     fprintf(stderr, "Usage: %s <width> <height> <ifile> <ofile> [options]\n", program);
 539     fprintf(stderr, "\t<width>  specifies the frame width\n");
 540     fprintf(stderr, "\t<height> specifies the frame height\n");
 541     fprintf(stderr, "\t<ifile>  specifies the I420/IYUV YUV file\n");
 542     fprintf(stderr, "\t<ofile>  specifies the encoded MPEG-2 file\n");
 543     fprintf(stderr, "where options include:\n");
 544     fprintf(stderr, "\t--cqp <QP>       const qp mode with specified <QP>\n");
 545     fprintf(stderr, "\t--fps <FPS>      specify the frame rate\n");
 546     fprintf(stderr, "\t--mode <MODE>    specify the mode 0 (I), 1 (I/P) and 2 (I/P/B)\n");
 547     fprintf(stderr, "\t--profile <PROFILE>      specify the profile 0(Simple), or 1(Main, default)\n");
 548     fprintf(stderr, "\t--level <LEVEL>  specify the level 0(Low), 1(Main, default) or 2(High)\n");
 549 }
 550
 551 void
 552 mpeg2_profile_level(struct mpeg2enc_context *ctx,
 553                     int profile,
 554                     int level)
 555 {
 556     int l = 2, p;
 557
 558     for (p = profile; p < 2; p++) {
 559         for (l = level; l < 3; l++) {
 560             if (ctx->width <= mpeg2_upper_samplings[p][l].samplers_per_line &&
 561                 ctx->height <= mpeg2_upper_samplings[p][l].line_per_frame &&
 562                 ctx->fps <= mpeg2_upper_samplings[p][l].frame_per_sec) {
 563
 564                 goto __find;
 565                 break;
 566             }
 567         }
 568     }
 569
 570     if (p == 2) {
 571         fprintf(stderr, "Warning: can't find a proper profile and level for the specified width/height/fps\n");
 572         p = 1;
 573         l = 2;
 574     }
 575
 576 __find:
 577     ctx->profile = mpeg2_va_profiles[p];
 578     ctx->level = l;
 579 }
 580
 581 static void
 582 parse_args(struct mpeg2enc_context *ctx, int argc, char **argv)
 583 {
 584     int c, tmp;
 585     int option_index = 0;
 586     long file_size;
 587     int profile = 1, level = 1;
 588
 589     static struct option long_options[] = {
 590         {"help",        no_argument,            0,      'h'},
 591         {"cqp",         required_argument,      0,      'c'},
 592         {"fps",         required_argument,      0,      'f'},
 593         {"mode",        required_argument,      0,      'm'},
 594         {"profile",     required_argument,      0,      'p'},
 595         {"level",       required_argument,      0,      'l'},
 596         { NULL,         0,                      NULL,   0 }
 597     };
 598
 599     if ((argc == 2 && strcmp(argv[1], "--help") == 0) ||
 600         (argc < 5))
 601         goto print_usage;
 602
 603     ctx->width = atoi(argv[1]);
 604     ctx->height = atoi(argv[2]);
 605
 606     if (ctx->width <= 0 || ctx->height <= 0) {
 607         fprintf(stderr, "<width> and <height> must be greater than 0\n");
 608         goto err_exit;
 609     }
 610
 611     ctx->ifp = fopen(argv[3], "rb");
 612
 613     if (ctx->ifp == NULL) {
 614         fprintf(stderr, "Can't open the input file\n");
 615         goto err_exit;
 616     }
 617
 618     fseek(ctx->ifp, 0l, SEEK_END);
 619     file_size = ftell(ctx->ifp);
 620     ctx->frame_size = ctx->width * ctx->height * 3 / 2;
 621
 622     if ((file_size < ctx->frame_size) ||
 623         (file_size % ctx->frame_size)) {
 624         fprintf(stderr, "The input file size %ld isn't a multiple of the frame size %d\n", file_size, ctx->frame_size);
 625         goto err_exit;
 626     }
 627
 628     ctx->num_pictures = file_size / ctx->frame_size;
 629     fseek(ctx->ifp, 0l, SEEK_SET);
 630
 631     ctx->ofp = fopen(argv[4], "wb");
 632
 633     if (ctx->ofp == NULL) {
 634         fprintf(stderr, "Can't create the output file\n");
 635         goto err_exit;
 636     }
 637
 638     opterr = 0;
 639     ctx->fps = 30;
 640     ctx->qp = 8;
 641     ctx->rate_control_mode = VA_RC_CQP;
 642     ctx->mode = MPEG2_MODE_IP;
 643     ctx->profile = VAProfileMPEG2Main;
 644     ctx->level = MPEG2_LEVEL_MAIN;
 645
 646     optind = 5;
 647
 648     while((c = getopt_long(argc, argv,
 649                            "",
 650                            long_options,
 651                            &option_index)) != -1) {
 652         switch(c) {
 653         case 'c':
 654             tmp = atoi(optarg);
 655
 656             /* only support q_scale_type = 0 */
 657             if (tmp > 62 || tmp < 2) {
 658                 fprintf(stderr, "Warning: QP must be in [2, 62]\n");
 659
 660                 if (tmp > 62)
 661                     tmp = 62;
 662
 663                 if (tmp < 2)
 664                     tmp = 2;
 665             }
 666
 667             ctx->qp = tmp & 0xFE;
 668             ctx->rate_control_mode = VA_RC_CQP;
 669
 670             break;
 671
 672         case 'f':
 673             tmp = atoi(optarg);
 674
 675             if (tmp <= 0)
 676                 fprintf(stderr, "Warning: FPS must be greater than 0\n");
 677             else
 678                 ctx->fps = tmp;
 679
 680             ctx->rate_control_mode = VA_RC_CBR;
 681
 682             break;
 683
 684         case 'm':
 685             tmp = atoi(optarg);
 686
 687             if (tmp < MPEG2_MODE_I || tmp > MPEG2_MODE_IPB)
 688                 fprintf(stderr, "Waning: MODE must be 0, 1, or 2\n");
 689             else
 690                 ctx->mode = tmp;
 691
 692             break;
 693
 694         case 'p':
 695             tmp = atoi(optarg);
 696
 697             if (tmp < 0 || tmp > 1)
 698                 fprintf(stderr, "Waning: PROFILE must be 0 or 1\n");
 699             else
 700                 profile = tmp;
 701
 702             break;
 703
 704         case 'l':
 705             tmp = atoi(optarg);
 706
 707             if (tmp < MPEG2_LEVEL_LOW || tmp > MPEG2_LEVEL_HIGH)
 708                 fprintf(stderr, "Waning: LEVEL must be 0, 1, or 2\n");
 709             else
 710                 level = tmp;
 711
 712             break;
 713
 714         case '?':
 715             fprintf(stderr, "Error: unkown command options\n");
 716
 717         case 'h':
 718             goto print_usage;
 719         }
 720     }
 721
 722     mpeg2_profile_level(ctx, profile, level);
 723
 724     return;
 725
 726 print_usage:
 727     usage(argv[0]);
 728 err_exit:
 729     mpeg2enc_exit(ctx, 1);
 730 }
 731
 732 /*
 733  * init
 734  */
 735 void
 736 mpeg2enc_init_sequence_parameter(struct mpeg2enc_context *ctx,
 737                                 VAEncSequenceParameterBufferMPEG2 *seq_param)
 738 {
 739     int profile = 4, level = 8;
 740
 741     switch (ctx->profile) {
 742     case VAProfileMPEG2Simple:
 743         profile = 5;
 744         break;
 745
 746     case VAProfileMPEG2Main:
 747         profile = 4;
 748         break;
 749
 750     default:
 751         assert(0);
 752         break;
 753     }
 754
 755     switch (ctx->level) {
 756     case MPEG2_LEVEL_LOW:
 757         level = 10;
 758         break;
 759
 760     case MPEG2_LEVEL_MAIN:
 761         level = 8;
 762         break;
 763
 764     case MPEG2_LEVEL_HIGH:
 765         level = 4;
 766         break;
 767
 768     default:
 769         assert(0);
 770         break;
 771     }
 772
 773     seq_param->intra_period = ctx->intra_period;
 774     seq_param->ip_period = ctx->ip_period;   /* FIXME: ??? */
 775     seq_param->picture_width = ctx->width;
 776     seq_param->picture_height = ctx->height;
 777
 778     if (ctx->bit_rate > 0)
 779         seq_param->bits_per_second = 1024 * ctx->bit_rate; /* use kbps as input */
 780     else
 781         seq_param->bits_per_second = 0x3FFFF * 400;
 782
 783     seq_param->frame_rate = ctx->fps;
 784     seq_param->aspect_ratio_information = 1;
 785     seq_param->vbv_buffer_size = 3; /* B = 16 * 1024 * vbv_buffer_size */
 786
 787     seq_param->sequence_extension.bits.profile_and_level_indication = profile << 4 | level;
 788     seq_param->sequence_extension.bits.progressive_sequence = 1; /* progressive frame-pictures */
 789     seq_param->sequence_extension.bits.chroma_format = CHROMA_FORMAT_420; /* 4:2:0 */
 790     seq_param->sequence_extension.bits.low_delay = 0; /* FIXME */
 791     seq_param->sequence_extension.bits.frame_rate_extension_n = 0;
 792     seq_param->sequence_extension.bits.frame_rate_extension_d = 0;
 793
 794     seq_param->gop_header.bits.time_code = (1 << 12); /* bit12: marker_bit */
 795     seq_param->gop_header.bits.closed_gop = 0;
 796     seq_param->gop_header.bits.broken_link = 0;
 797 }
 798
 799 static void
 800 mpeg2enc_init_picture_parameter(struct mpeg2enc_context *ctx,
 801                                VAEncPictureParameterBufferMPEG2 *pic_param)
 802 {
 803     pic_param->forward_reference_picture = VA_INVALID_ID;
 804     pic_param->backward_reference_picture = VA_INVALID_ID;
 805     pic_param->reconstructed_picture = VA_INVALID_ID;
 806     pic_param->coded_buf = VA_INVALID_ID;
 807     pic_param->picture_type = VAEncPictureTypeIntra;
 808
 809     pic_param->temporal_reference = 0;
 810     pic_param->f_code[0][0] = 0xf;
 811     pic_param->f_code[0][1] = 0xf;
 812     pic_param->f_code[1][0] = 0xf;
 813     pic_param->f_code[1][1] = 0xf;
 814
 815     pic_param->picture_coding_extension.bits.intra_dc_precision = 0; /* 8bits */
 816     pic_param->picture_coding_extension.bits.picture_structure = 3; /* frame picture */
 817     pic_param->picture_coding_extension.bits.top_field_first = 0;
 818     pic_param->picture_coding_extension.bits.frame_pred_frame_dct = 1; /* FIXME */
 819     pic_param->picture_coding_extension.bits.concealment_motion_vectors = 0;
 820     pic_param->picture_coding_extension.bits.q_scale_type = 0;
 821     pic_param->picture_coding_extension.bits.intra_vlc_format = 0;
 822     pic_param->picture_coding_extension.bits.alternate_scan = 0;
 823     pic_param->picture_coding_extension.bits.repeat_first_field = 0;
 824     pic_param->picture_coding_extension.bits.progressive_frame = 1;
 825     pic_param->picture_coding_extension.bits.composite_display_flag = 0;
 826
 827     pic_param->user_data_length = 0;
 828 }
 829
 830 static void
 831 mpeg2enc_alloc_va_resources(struct mpeg2enc_context *ctx)
 832 {
 833     VAEntrypoint *entrypoint_list;
 834     VAConfigAttrib attrib_list[2];
 835     VAStatus va_status;
 836     int max_entrypoints, num_entrypoints, entrypoint;
 837     int major_ver, minor_ver;
 838
 839     ctx->va_dpy = va_open_display();
 840     va_status = vaInitialize(ctx->va_dpy,
 841                              &major_ver,
 842                              &minor_ver);
 843     CHECK_VASTATUS(va_status, "vaInitialize");
 844
 845     max_entrypoints = vaMaxNumEntrypoints(ctx->va_dpy);
 846     entrypoint_list = malloc(max_entrypoints * sizeof(VAEntrypoint));
 847     vaQueryConfigEntrypoints(ctx->va_dpy,
 848                              ctx->profile,
 849                              entrypoint_list,
 850                              &num_entrypoints);
 851
 852     for (entrypoint = 0; entrypoint < num_entrypoints; entrypoint++) {
 853         if (entrypoint_list[entrypoint] == VAEntrypointEncSlice)
 854             break;
 855     }
 856
 857     free(entrypoint_list);
 858
 859     if (entrypoint == num_entrypoints) {
 860         /* not find Slice entry point */
 861         assert(0);
 862     }
 863
 864     /* find out the format for the render target, and rate control mode */
 865     attrib_list[0].type = VAConfigAttribRTFormat;
 866     attrib_list[1].type = VAConfigAttribRateControl;
 867     vaGetConfigAttributes(ctx->va_dpy,
 868                           ctx->profile,
 869                           VAEntrypointEncSlice,
 870                           &attrib_list[0],
 871                           2);
 872
 873     if ((attrib_list[0].value & VA_RT_FORMAT_YUV420) == 0) {
 874         /* not find desired YUV420 RT format */
 875         assert(0);
 876     }
 877
 878     if ((attrib_list[1].value & ctx->rate_control_mode) == 0) {
 879         /* Can't find matched RC mode */
 880         fprintf(stderr, "RC mode %d isn't found, exit\n", ctx->rate_control_mode);
 881         assert(0);
 882     }
 883
 884     attrib_list[0].value = VA_RT_FORMAT_YUV420; /* set to desired RT format */
 885     attrib_list[1].value = ctx->rate_control_mode; /* set to desired RC mode */
 886
 887     va_status = vaCreateConfig(ctx->va_dpy,
 888                                ctx->profile,
 889                                VAEntrypointEncSlice,
 890                                attrib_list,
 891                                2,
 892                                &ctx->config_id);
 893     CHECK_VASTATUS(va_status, "vaCreateConfig");
 894
 895     /* Create a context for this decode pipe */
 896     va_status = vaCreateContext(ctx->va_dpy,
 897                                 ctx->config_id,
 898                                 ctx->width,
 899                                 ctx->height,
 900                                 VA_PROGRESSIVE,
 901                                 0,
 902                                 0,
 903                                 &ctx->context_id);
 904     CHECK_VASTATUS(va_status, "vaCreateContext");
 905
 906     va_status = vaCreateSurfaces(ctx->va_dpy,
 907                                  VA_RT_FORMAT_YUV420,
 908                                  ctx->width,
 909                                  ctx->height,
 910                                  surface_ids,
 911                                  SID_NUMBER,
 912                                  NULL,
 913                                  0);
 914     CHECK_VASTATUS(va_status, "vaCreateSurfaces");
 915 }
 916
 917 static void
 918 mpeg2enc_init(struct mpeg2enc_context *ctx)
 919 {
 920     int i;
 921
 922     ctx->frame_data_buffer = (unsigned char *)malloc(ctx->frame_size);
 923     ctx->seq_param_buf_id = VA_INVALID_ID;
 924     ctx->pic_param_buf_id = VA_INVALID_ID;
 925     ctx->packed_seq_header_param_buf_id = VA_INVALID_ID;
 926     ctx->packed_seq_buf_id = VA_INVALID_ID;
 927     ctx->packed_pic_header_param_buf_id = VA_INVALID_ID;
 928     ctx->packed_pic_buf_id = VA_INVALID_ID;
 929     ctx->codedbuf_buf_id = VA_INVALID_ID;
 930     ctx->codedbuf_i_size = ctx->frame_size;
 931     ctx->codedbuf_pb_size = 0;
 932     ctx->next_display_order = 0;
 933     ctx->next_type = VAEncPictureTypeIntra;
 934
 935     if (ctx->mode == MPEG2_MODE_I) {
 936         ctx->intra_period = 1;
 937         ctx->ip_period = 0;
 938     } else if (ctx->mode == MPEG2_MODE_IP) {
 939         ctx->intra_period = 16;
 940         ctx->ip_period = 0;
 941     } else {
 942         ctx->intra_period = 16;
 943         ctx->ip_period = 2;
 944     }
 945
 946     ctx->next_bframes = ctx->ip_period;
 947
 948     ctx->new_sequence = 1;
 949     ctx->new_gop_header = 1;
 950     ctx->gop_header_in_display_order = 0;
 951
 952     ctx->bit_rate = -1;
 953
 954     for (i = 0; i < MAX_SLICES; i++) {
 955         ctx->slice_param_buf_id[i] = VA_INVALID_ID;
 956     }
 957
 958     mpeg2enc_init_sequence_parameter(ctx, &ctx->seq_param);
 959     mpeg2enc_init_picture_parameter(ctx, &ctx->pic_param);
 960     mpeg2enc_alloc_va_resources(ctx);
 961
 962     /* thread */
 963     ctx->current_input_surface = SID_INPUT_PICTURE_0;
 964     ctx->current_upload_surface = SID_INPUT_PICTURE_1;
 965     ctx->upload_thread_value = pthread_create(&ctx->upload_thread_id,
 966                                               NULL,
 967                                               upload_yuv_to_surface,
 968                                               ctx);
 969 }
 970
 971 static int
 972 mpeg2enc_time_code(VAEncSequenceParameterBufferMPEG2 *seq_param,
 973                    int num_frames)
 974 {
 975     int fps = (int)(seq_param->frame_rate + 0.5);
 976     int time_code = 0;
 977     int time_code_pictures, time_code_seconds, time_code_minutes, time_code_hours;
 978     int drop_frame_flag = 0;
 979
 980     assert(fps <= 60);
 981
 982     time_code_seconds = num_frames / fps;
 983     time_code_pictures = num_frames % fps;
 984     time_code |= time_code_pictures;
 985
 986     time_code_minutes = time_code_seconds / 60;
 987     time_code_seconds = time_code_seconds % 60;
 988     time_code |= (time_code_seconds << 6);
 989
 990     time_code_hours = time_code_minutes / 60;
 991     time_code_minutes = time_code_minutes % 60;
 992
 993     time_code |= (1 << 12);     /* marker_bit */
 994     time_code |= (time_code_minutes << 13);
 995
 996     time_code_hours = time_code_hours % 24;
 997     time_code |= (time_code_hours << 19);
 998
 999     time_code |= (drop_frame_flag << 24);
1000
1001     return time_code;
1002 }
1003
1004 /*
1005  * run
1006  */
1007 static void
1008 mpeg2enc_update_sequence_parameter(struct mpeg2enc_context *ctx,
1009                                    VAEncPictureType picture_type,
1010                                    int coded_order,
1011                                    int display_order)
1012 {
1013     VAEncSequenceParameterBufferMPEG2 *seq_param = &ctx->seq_param;
1014
1015     /* update the time_code info for the new GOP */
1016     if (ctx->new_gop_header) {
1017         seq_param->gop_header.bits.time_code = mpeg2enc_time_code(seq_param, display_order);
1018     }
1019 }
1020
1021 static void
1022 mpeg2enc_update_picture_parameter(struct mpeg2enc_context *ctx,
1023                                   VAEncPictureType picture_type,
1024                                   int coded_order,
1025                                   int display_order)
1026 {
1027     VAEncPictureParameterBufferMPEG2 *pic_param = &ctx->pic_param;
1028
1029     pic_param->picture_type = picture_type;
1030     pic_param->temporal_reference = (display_order - ctx->gop_header_in_display_order) & 0x3FF;
1031     pic_param->reconstructed_picture = surface_ids[SID_RECON_PICTURE];
1032     pic_param->forward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L0];
1033     pic_param->backward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L1];
1034
1035     if (pic_param->picture_type == VAEncPictureTypeIntra) {
1036         pic_param->f_code[0][0] = 0xf;
1037         pic_param->f_code[0][1] = 0xf;
1038         pic_param->f_code[1][0] = 0xf;
1039         pic_param->f_code[1][1] = 0xf;
1040         pic_param->forward_reference_picture = VA_INVALID_SURFACE;
1041         pic_param->backward_reference_picture = VA_INVALID_SURFACE;
1042
1043     } else if (pic_param->picture_type == VAEncPictureTypePredictive) {
1044         pic_param->f_code[0][0] = 0x1;
1045         pic_param->f_code[0][1] = 0x1;
1046         pic_param->f_code[1][0] = 0xf;
1047         pic_param->f_code[1][1] = 0xf;
1048         pic_param->forward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L0];
1049         pic_param->backward_reference_picture = VA_INVALID_SURFACE;
1050     } else if (pic_param->picture_type == VAEncPictureTypeBidirectional) {
1051         pic_param->f_code[0][0] = 0x1;
1052         pic_param->f_code[0][1] = 0x1;
1053         pic_param->f_code[1][0] = 0x1;
1054         pic_param->f_code[1][1] = 0x1;
1055         pic_param->forward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L0];
1056         pic_param->backward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L1];
1057     } else {
1058         assert(0);
1059     }
1060 }
1061
1062 static void
1063 mpeg2enc_update_picture_parameter_buffer(struct mpeg2enc_context *ctx,
1064                                          VAEncPictureType picture_type,
1065                                          int coded_order,
1066                                          int display_order)
1067 {
1068     VAEncPictureParameterBufferMPEG2 *pic_param = &ctx->pic_param;
1069     VAStatus va_status;
1070
1071     /* update the coded buffer id */
1072     pic_param->coded_buf = ctx->codedbuf_buf_id;
1073     va_status = vaCreateBuffer(ctx->va_dpy,
1074                                ctx->context_id,
1075                                VAEncPictureParameterBufferType,
1076                                sizeof(*pic_param),
1077                                1,
1078                                pic_param,
1079                                &ctx->pic_param_buf_id);
1080     CHECK_VASTATUS(va_status, "vaCreateBuffer");
1081 }
1082
1083 static void
1084 mpeg2enc_update_slice_parameter(struct mpeg2enc_context *ctx, VAEncPictureType picture_type)
1085 {
1086     VAEncSequenceParameterBufferMPEG2 *seq_param;
1087     VAEncPictureParameterBufferMPEG2 *pic_param;
1088     VAEncSliceParameterBufferMPEG2 *slice_param;
1089     VAStatus va_status;
1090     int i, width_in_mbs, height_in_mbs;
1091
1092     pic_param = &ctx->pic_param;
1093     assert(pic_param->picture_coding_extension.bits.q_scale_type == 0);
1094
1095     seq_param = &ctx->seq_param;
1096     width_in_mbs = (seq_param->picture_width + 15) / 16;
1097     height_in_mbs = (seq_param->picture_height + 15) / 16;
1098     ctx->num_slice_groups = 1;
1099
1100     for (i = 0; i < height_in_mbs; i++) {
1101         slice_param = &ctx->slice_param[i];
1102         slice_param->macroblock_address = i * width_in_mbs;
1103         slice_param->num_macroblocks = width_in_mbs;
1104         slice_param->is_intra_slice = (picture_type == VAEncPictureTypeIntra);
1105         slice_param->quantiser_scale_code = ctx->qp / 2;
1106     }
1107
1108     va_status = vaCreateBuffer(ctx->va_dpy,
1109                                ctx->context_id,
1110                                VAEncSliceParameterBufferType,
1111                                sizeof(*slice_param),
1112                                height_in_mbs,
1113                                ctx->slice_param,
1114                                ctx->slice_param_buf_id);
1115     CHECK_VASTATUS(va_status, "vaCreateBuffer");;
1116 }
1117
1118 static int
1119 begin_picture(struct mpeg2enc_context *ctx,
1120               int coded_order,
1121               int display_order,
1122               VAEncPictureType picture_type)
1123 {
1124     VAStatus va_status;
1125     int tmp;
1126     VAEncPackedHeaderParameterBuffer packed_header_param_buffer;
1127     unsigned int length_in_bits;
1128     unsigned char *packed_seq_buffer = NULL, *packed_pic_buffer = NULL;
1129
1130     if (ctx->upload_thread_value != 0) {
1131         fprintf(stderr, "FATAL error!!!\n");
1132         exit(1);
1133     }
1134
1135     pthread_join(ctx->upload_thread_id, NULL);
1136
1137     ctx->upload_thread_value = -1;
1138     tmp = ctx->current_input_surface;
1139     ctx->current_input_surface = ctx->current_upload_surface;
1140     ctx->current_upload_surface = tmp;
1141
1142     mpeg2enc_update_sequence_parameter(ctx, picture_type, coded_order, display_order);
1143     mpeg2enc_update_picture_parameter(ctx, picture_type, coded_order, display_order);
1144
1145     if (ctx->new_sequence || ctx->new_gop_header) {
1146         assert(picture_type == VAEncPictureTypeIntra);
1147         length_in_bits = build_packed_seq_buffer(ctx, &ctx->seq_param, &packed_seq_buffer);
1148         packed_header_param_buffer.type = VAEncPackedHeaderMPEG2_SPS;
1149         packed_header_param_buffer.has_emulation_bytes = 0;
1150         packed_header_param_buffer.bit_length = length_in_bits;
1151         va_status = vaCreateBuffer(ctx->va_dpy,
1152                                    ctx->context_id,
1153                                    VAEncPackedHeaderParameterBufferType,
1154                                    sizeof(packed_header_param_buffer), 1, &packed_header_param_buffer,
1155                                    &ctx->packed_seq_header_param_buf_id);
1156         CHECK_VASTATUS(va_status,"vaCreateBuffer");
1157
1158         va_status = vaCreateBuffer(ctx->va_dpy,
1159                                    ctx->context_id,
1160                                    VAEncPackedHeaderDataBufferType,
1161                                    (length_in_bits + 7) / 8, 1, packed_seq_buffer,
1162                                    &ctx->packed_seq_buf_id);
1163         CHECK_VASTATUS(va_status,"vaCreateBuffer");
1164
1165         free(packed_seq_buffer);
1166     }
1167
1168     length_in_bits = build_packed_pic_buffer(&ctx->seq_param, &ctx->pic_param, &packed_pic_buffer);
1169     packed_header_param_buffer.type = VAEncPackedHeaderMPEG2_PPS;
1170     packed_header_param_buffer.has_emulation_bytes = 0;
1171     packed_header_param_buffer.bit_length = length_in_bits;
1172
1173     va_status = vaCreateBuffer(ctx->va_dpy,
1174                                ctx->context_id,
1175                                VAEncPackedHeaderParameterBufferType,
1176                                sizeof(packed_header_param_buffer), 1, &packed_header_param_buffer,
1177                                &ctx->packed_pic_header_param_buf_id);
1178     CHECK_VASTATUS(va_status,"vaCreateBuffer");
1179
1180     va_status = vaCreateBuffer(ctx->va_dpy,
1181                                ctx->context_id,
1182                                VAEncPackedHeaderDataBufferType,
1183                                (length_in_bits + 7) / 8, 1, packed_pic_buffer,
1184                                &ctx->packed_pic_buf_id);
1185     CHECK_VASTATUS(va_status,"vaCreateBuffer");
1186
1187     free(packed_pic_buffer);
1188
1189     /* sequence parameter set */
1190     VAEncSequenceParameterBufferMPEG2 *seq_param = &ctx->seq_param;
1191     va_status = vaCreateBuffer(ctx->va_dpy,
1192                                ctx->context_id,
1193                                VAEncSequenceParameterBufferType,
1194                                sizeof(*seq_param),
1195                                1,
1196                                seq_param,
1197                                &ctx->seq_param_buf_id);
1198     CHECK_VASTATUS(va_status,"vaCreateBuffer");;
1199
1200     /* slice parameter */
1201     mpeg2enc_update_slice_parameter(ctx, picture_type);
1202
1203     return 0;
1204 }
1205
1206 static int
1207 mpeg2enc_render_picture(struct mpeg2enc_context *ctx)
1208 {
1209     VAStatus va_status;
1210     VABufferID va_buffers[16];
1211     unsigned int num_va_buffers = 0;
1212
1213     va_buffers[num_va_buffers++] = ctx->seq_param_buf_id;
1214     va_buffers[num_va_buffers++] = ctx->pic_param_buf_id;
1215
1216     if (ctx->packed_seq_header_param_buf_id != VA_INVALID_ID)
1217         va_buffers[num_va_buffers++] = ctx->packed_seq_header_param_buf_id;
1218
1219     if (ctx->packed_seq_buf_id != VA_INVALID_ID)
1220         va_buffers[num_va_buffers++] = ctx->packed_seq_buf_id;
1221
1222     if (ctx->packed_pic_header_param_buf_id != VA_INVALID_ID)
1223         va_buffers[num_va_buffers++] = ctx->packed_pic_header_param_buf_id;
1224
1225     if (ctx->packed_pic_buf_id != VA_INVALID_ID)
1226         va_buffers[num_va_buffers++] = ctx->packed_pic_buf_id;
1227
1228     va_status = vaBeginPicture(ctx->va_dpy,
1229                                ctx->context_id,
1230                                surface_ids[ctx->current_input_surface]);
1231     CHECK_VASTATUS(va_status,"vaBeginPicture");
1232
1233     va_status = vaRenderPicture(ctx->va_dpy,
1234                                 ctx->context_id,
1235                                 va_buffers,
1236                                 num_va_buffers);
1237     CHECK_VASTATUS(va_status,"vaRenderPicture");
1238
1239     va_status = vaRenderPicture(ctx->va_dpy,
1240                                 ctx->context_id,
1241                                 &ctx->slice_param_buf_id[0],
1242                                 ctx->num_slice_groups);
1243     CHECK_VASTATUS(va_status,"vaRenderPicture");
1244
1245     va_status = vaEndPicture(ctx->va_dpy, ctx->context_id);
1246     CHECK_VASTATUS(va_status,"vaEndPicture");
1247
1248     return 0;
1249 }
1250
1251 static int
1252 mpeg2enc_destroy_buffers(struct mpeg2enc_context *ctx, VABufferID *va_buffers, unsigned int num_va_buffers)
1253 {
1254     VAStatus va_status;
1255     unsigned int i;
1256
1257     for (i = 0; i < num_va_buffers; i++) {
1258         if (va_buffers[i] != VA_INVALID_ID) {
1259             va_status = vaDestroyBuffer(ctx->va_dpy, va_buffers[i]);
1260             CHECK_VASTATUS(va_status,"vaDestroyBuffer");
1261             va_buffers[i] = VA_INVALID_ID;
1262         }
1263     }
1264
1265     return 0;
1266 }
1267
1268 static void
1269 end_picture(struct mpeg2enc_context *ctx, VAEncPictureType picture_type, int next_is_bpic)
1270 {
1271     VABufferID tempID;
1272
1273     /* Prepare for next picture */
1274     tempID = surface_ids[SID_RECON_PICTURE];
1275
1276     if (picture_type != VAEncPictureTypeBidirectional) {
1277         if (next_is_bpic) {
1278             surface_ids[SID_RECON_PICTURE] = surface_ids[SID_REFERENCE_PICTURE_L1];
1279             surface_ids[SID_REFERENCE_PICTURE_L1] = tempID;
1280         } else {
1281             surface_ids[SID_RECON_PICTURE] = surface_ids[SID_REFERENCE_PICTURE_L0];
1282             surface_ids[SID_REFERENCE_PICTURE_L0] = tempID;
1283         }
1284     } else {
1285         if (!next_is_bpic) {
1286             surface_ids[SID_RECON_PICTURE] = surface_ids[SID_REFERENCE_PICTURE_L0];
1287             surface_ids[SID_REFERENCE_PICTURE_L0] = surface_ids[SID_REFERENCE_PICTURE_L1];
1288             surface_ids[SID_REFERENCE_PICTURE_L1] = tempID;
1289         }
1290     }
1291
1292     mpeg2enc_destroy_buffers(ctx, &ctx->seq_param_buf_id, 1);
1293     mpeg2enc_destroy_buffers(ctx, &ctx->pic_param_buf_id, 1);
1294     mpeg2enc_destroy_buffers(ctx, &ctx->packed_seq_header_param_buf_id, 1);
1295     mpeg2enc_destroy_buffers(ctx, &ctx->packed_seq_buf_id, 1);
1296     mpeg2enc_destroy_buffers(ctx, &ctx->packed_pic_header_param_buf_id, 1);
1297     mpeg2enc_destroy_buffers(ctx, &ctx->packed_pic_buf_id, 1);
1298     mpeg2enc_destroy_buffers(ctx, &ctx->slice_param_buf_id[0], ctx->num_slice_groups);
1299     mpeg2enc_destroy_buffers(ctx, &ctx->codedbuf_buf_id, 1);
1300     memset(ctx->slice_param, 0, sizeof(ctx->slice_param));
1301     ctx->num_slice_groups = 0;
1302 }
1303
1304 static int
1305 store_coded_buffer(struct mpeg2enc_context *ctx, VAEncPictureType picture_type)
1306 {
1307     VACodedBufferSegment *coded_buffer_segment;
1308     unsigned char *coded_mem;
1309     int slice_data_length;
1310     VAStatus va_status;
1311     VASurfaceStatus surface_status;
1312     size_t w_items;
1313
1314     va_status = vaSyncSurface(ctx->va_dpy, surface_ids[ctx->current_input_surface]);
1315     CHECK_VASTATUS(va_status,"vaSyncSurface");
1316
1317     surface_status = 0;
1318     va_status = vaQuerySurfaceStatus(ctx->va_dpy, surface_ids[ctx->current_input_surface], &surface_status);
1319     CHECK_VASTATUS(va_status,"vaQuerySurfaceStatus");
1320
1321     va_status = vaMapBuffer(ctx->va_dpy, ctx->codedbuf_buf_id, (void **)(&coded_buffer_segment));
1322     CHECK_VASTATUS(va_status,"vaMapBuffer");
1323     coded_mem = coded_buffer_segment->buf;
1324
1325     if (coded_buffer_segment->status & VA_CODED_BUF_STATUS_SLICE_OVERFLOW_MASK) {
1326         if (picture_type == VAEncPictureTypeIntra)
1327             ctx->codedbuf_i_size *= 2;
1328         else
1329             ctx->codedbuf_pb_size *= 2;
1330
1331         vaUnmapBuffer(ctx->va_dpy, ctx->codedbuf_buf_id);
1332         return -1;
1333     }
1334
1335     slice_data_length = coded_buffer_segment->size;
1336
1337     do {
1338         w_items = fwrite(coded_mem, slice_data_length, 1, ctx->ofp);
1339     } while (w_items != 1);
1340
1341     if (picture_type == VAEncPictureTypeIntra) {
1342         if (ctx->codedbuf_i_size > slice_data_length * 3 / 2) {
1343             ctx->codedbuf_i_size = slice_data_length * 3 / 2;
1344         }
1345
1346         if (ctx->codedbuf_pb_size < slice_data_length) {
1347             ctx->codedbuf_pb_size = slice_data_length;
1348         }
1349     } else {
1350         if (ctx->codedbuf_pb_size > slice_data_length * 3 / 2) {
1351             ctx->codedbuf_pb_size = slice_data_length * 3 / 2;
1352         }
1353     }
1354
1355     vaUnmapBuffer(ctx->va_dpy, ctx->codedbuf_buf_id);
1356
1357     return 0;
1358 }
1359
1360 static void
1361 encode_picture(struct mpeg2enc_context *ctx,
1362                int coded_order,
1363                int display_order,
1364                VAEncPictureType picture_type,
1365                int next_is_bpic,
1366                int next_display_order)
1367 {
1368     VAStatus va_status;
1369     int ret = 0, codedbuf_size;
1370
1371     begin_picture(ctx, coded_order, display_order, picture_type);
1372
1373     if (1) {
1374         /* upload YUV data to VA surface for next frame */
1375         if (next_display_order >= ctx->num_pictures)
1376             next_display_order = ctx->num_pictures - 1;
1377
1378         fseek(ctx->ifp, ctx->frame_size * next_display_order, SEEK_SET);
1379         ctx->upload_thread_value = pthread_create(&ctx->upload_thread_id,
1380                                                   NULL,
1381                                                   upload_yuv_to_surface,
1382                                                   ctx);
1383     }
1384
1385     do {
1386         mpeg2enc_destroy_buffers(ctx, &ctx->codedbuf_buf_id, 1);
1387         mpeg2enc_destroy_buffers(ctx, &ctx->pic_param_buf_id, 1);
1388
1389
1390         if (VAEncPictureTypeIntra == picture_type) {
1391             codedbuf_size = ctx->codedbuf_i_size;
1392         } else {
1393             codedbuf_size = ctx->codedbuf_pb_size;
1394         }
1395
1396         /* coded buffer */
1397         va_status = vaCreateBuffer(ctx->va_dpy,
1398                                    ctx->context_id,
1399                                    VAEncCodedBufferType,
1400                                    codedbuf_size, 1, NULL,
1401                                    &ctx->codedbuf_buf_id);
1402         CHECK_VASTATUS(va_status,"vaCreateBuffer");
1403
1404         /* picture parameter set */
1405         mpeg2enc_update_picture_parameter_buffer(ctx, picture_type, coded_order, display_order);
1406
1407         mpeg2enc_render_picture(ctx);
1408
1409         ret = store_coded_buffer(ctx, picture_type);
1410     } while (ret);
1411
1412     end_picture(ctx, picture_type, next_is_bpic);
1413 }
1414
1415 static void
1416 update_next_frame_info(struct mpeg2enc_context *ctx,
1417                        VAEncPictureType curr_type,
1418                        int curr_coded_order,
1419                        int curr_display_order)
1420 {
1421     if (((curr_coded_order + 1) % ctx->intra_period) == 0) {
1422         ctx->next_type = VAEncPictureTypeIntra;
1423         ctx->next_display_order = curr_coded_order + 1;
1424
1425         return;
1426     }
1427
1428     if (curr_type == VAEncPictureTypeIntra) {
1429         assert(curr_display_order == curr_coded_order);
1430         ctx->next_type = VAEncPictureTypePredictive;
1431         ctx->next_bframes = ctx->ip_period;
1432         ctx->next_display_order = curr_display_order + ctx->next_bframes + 1;
1433     } else if (curr_type == VAEncPictureTypePredictive) {
1434         if (ctx->ip_period == 0) {
1435             assert(curr_display_order == curr_coded_order);
1436             ctx->next_type = VAEncPictureTypePredictive;
1437             ctx->next_display_order = curr_display_order + 1;
1438         } else {
1439             ctx->next_type = VAEncPictureTypeBidirectional;
1440             ctx->next_display_order = curr_display_order - ctx->next_bframes;
1441             ctx->next_bframes--;
1442         }
1443     } else if (curr_type == VAEncPictureTypeBidirectional) {
1444         if (ctx->next_bframes == 0) {
1445             ctx->next_type = VAEncPictureTypePredictive;
1446             ctx->next_bframes = ctx->ip_period;
1447             ctx->next_display_order = curr_display_order + ctx->next_bframes + 2;
1448         } else {
1449             ctx->next_type = VAEncPictureTypeBidirectional;
1450             ctx->next_display_order = curr_display_order + 1;
1451             ctx->next_bframes--;
1452         }
1453     }
1454
1455     if (ctx->next_display_order >= ctx->num_pictures) {
1456         int rtmp = ctx->next_display_order - (ctx->num_pictures - 1);
1457         ctx->next_display_order = ctx->num_pictures - 1;
1458         ctx->next_bframes -= rtmp;
1459     }
1460 }
1461
1462 static void
1463 mpeg2enc_run(struct mpeg2enc_context *ctx)
1464 {
1465     int display_order = 0, coded_order = 0;
1466     VAEncPictureType type;
1467
1468     ctx->new_sequence = 1;
1469     ctx->new_gop_header = 1;
1470     ctx->gop_header_in_display_order = display_order;
1471
1472     while (coded_order < ctx->num_pictures) {
1473         type = ctx->next_type;
1474         display_order = ctx->next_display_order;
1475         /* follow the IPBxxBPBxxB mode */
1476         update_next_frame_info(ctx, type, coded_order, display_order);
1477         encode_picture(ctx,
1478                        coded_order,
1479                        display_order,
1480                        type,
1481                        ctx->next_type == VAEncPictureTypeBidirectional,
1482                        ctx->next_display_order);
1483
1484         /* update gop_header */
1485         ctx->new_sequence = 0;
1486         ctx->new_gop_header = ctx->next_type == VAEncPictureTypeIntra;
1487
1488         if (ctx->new_gop_header)
1489             ctx->gop_header_in_display_order += ctx->intra_period;
1490
1491         coded_order++;
1492
1493         fprintf(stderr, "\r %d/%d ...", coded_order, ctx->num_pictures);
1494         fflush(stdout);
1495     }
1496 }
1497
1498 /*
1499  * end
1500  */
1501 static void
1502 mpeg2enc_release_va_resources(struct mpeg2enc_context *ctx)
1503 {
1504     vaDestroySurfaces(ctx->va_dpy, surface_ids, SID_NUMBER);
1505     vaDestroyContext(ctx->va_dpy, ctx->context_id);
1506     vaDestroyConfig(ctx->va_dpy, ctx->config_id);
1507     vaTerminate(ctx->va_dpy);
1508     va_close_display(ctx->va_dpy);
1509 }
1510
1511 static void
1512 mpeg2enc_end(struct mpeg2enc_context *ctx)
1513 {
1514     pthread_join(ctx->upload_thread_id, NULL);
1515     mpeg2enc_release_va_resources(ctx);
1516 }
1517
1518 int
1519 main(int argc, char *argv[])
1520 {
1521     struct mpeg2enc_context ctx;
1522     struct timeval tpstart, tpend;
1523     float timeuse;
1524
1525     gettimeofday(&tpstart, NULL);
1526
1527     memset(&ctx, 0, sizeof(ctx));
1528     parse_args(&ctx, argc, argv);
1529     mpeg2enc_init(&ctx);
1530     mpeg2enc_run(&ctx);
1531     mpeg2enc_end(&ctx);
1532
1533     gettimeofday(&tpend, NULL);
1534     timeuse = 1000000 * (tpend.tv_sec - tpstart.tv_sec) + tpend.tv_usec - tpstart.tv_usec;
1535     timeuse /= 1000000;
1536     fprintf(stderr, "\ndone!\n");
1537     fprintf(stderr, "encode %d frames in %f secondes, FPS is %.1f\n", ctx.num_pictures, timeuse, ctx.num_pictures / timeuse);
1538
1539     mpeg2enc_exit(&ctx, 0);
1540
1541     return 0;
1542 }