Modify set refernce example to work with HBD path

The default configuration was to use HBD path, which was the reason that
the example test failed. Modified the example code so that it supported
the internal HBD path.

BUG=aomedia:1842

Change-Id: Ie76ba93059ae1fb705e0e8989e3f5242e55fa931
diff --git a/examples/aom_cx_set_ref.c b/examples/aom_cx_set_ref.c
index b40885c..923caa8 100644
--- a/examples/aom_cx_set_ref.c
+++ b/examples/aom_cx_set_ref.c
@@ -184,11 +184,16 @@
   aom_codec_enc_cfg_t cfg;
   unsigned int frame_in = 0;
   aom_image_t raw;
+  aom_image_t raw_shift;
   aom_image_t ext_ref;
   aom_codec_err_t res;
   AvxVideoInfo info;
   AvxVideoWriter *writer = NULL;
   const AvxInterface *encoder = NULL;
+  int flags = 0;
+  int allocated_raw_shift = 0;
+  aom_img_fmt_t raw_fmt = AOM_IMG_FMT_I420;
+  aom_img_fmt_t ref_fmt = AOM_IMG_FMT_I420;
 
   // Test encoder/decoder mismatch.
   int test_decode = 1;
@@ -254,12 +259,15 @@
     die("Invalid frame size: %dx%d", info.frame_width, info.frame_height);
   }
 
-  if (!aom_img_alloc_with_border(&raw, AOM_IMG_FMT_I420, info.frame_width,
-                                 info.frame_height, 1, 8, 0)) {
+  // In this test, the bit depth of input video is 8-bit, and the input format
+  // is AOM_IMG_FMT_I420.
+  if (!aom_img_alloc(&raw, raw_fmt, info.frame_width, info.frame_height, 32)) {
     die("Failed to allocate image.");
   }
+
+  if (!CONFIG_LOWBITDEPTH) ref_fmt |= AOM_IMG_FMT_HIGHBITDEPTH;
   // Allocate memory with the border so that it can be used as a reference.
-  if (!aom_img_alloc_with_border(&ext_ref, AOM_IMG_FMT_I420, info.frame_width,
+  if (!aom_img_alloc_with_border(&ext_ref, ref_fmt, info.frame_width,
                                  info.frame_height, 32, 8,
                                  AOM_BORDER_IN_PIXELS)) {
     die("Failed to allocate image.");
@@ -276,6 +284,11 @@
   cfg.g_timebase.den = info.time_base.denominator;
   cfg.rc_target_bitrate = bitrate;
   cfg.g_lag_in_frames = 3;
+  cfg.g_bit_depth = AOM_BITS_8;
+
+  flags |= (cfg.g_bit_depth > AOM_BITS_8 || !CONFIG_LOWBITDEPTH)
+               ? AOM_CODEC_USE_HIGHBITDEPTH
+               : 0;
 
   writer = aom_video_writer_open(outfile_arg, kContainerIVF, &info);
   if (!writer) die("Failed to open %s for writing.", outfile_arg);
@@ -283,7 +296,7 @@
   if (!(infile = fopen(infile_arg, "rb")))
     die("Failed to open %s for reading.", infile_arg);
 
-  if (aom_codec_enc_init(&ecodec, encoder->codec_interface(), &cfg, 0))
+  if (aom_codec_enc_init(&ecodec, encoder->codec_interface(), &cfg, flags))
     die_codec(&ecodec, "Failed to initialize encoder");
 
   // Disable alt_ref.
@@ -299,6 +312,22 @@
   // Encode frames.
   while (aom_img_read(&raw, infile)) {
     if (limit && frame_in >= limit) break;
+    aom_image_t *frame_to_encode;
+
+    if (!CONFIG_LOWBITDEPTH) {
+      // Need to allocate larger buffer to use hbd internal.
+      int input_shift = 0;
+      if (!allocated_raw_shift) {
+        aom_img_alloc(&raw_shift, raw_fmt | AOM_IMG_FMT_HIGHBITDEPTH,
+                      info.frame_width, info.frame_height, 32);
+        allocated_raw_shift = 1;
+      }
+      aom_img_upshift(&raw_shift, &raw, input_shift);
+      frame_to_encode = &raw_shift;
+    } else {
+      frame_to_encode = &raw;
+    }
+
     if (update_frame_num > 1 && frame_out + 1 == update_frame_num) {
       av1_ref_frame_t ref;
       ref.idx = 0;
@@ -318,8 +347,8 @@
       }
     }
 
-    encode_frame(&ecodec, &raw, frame_in, writer, test_decode, &dcodec,
-                 &frame_out, &mismatch_seen, &ext_ref);
+    encode_frame(&ecodec, frame_to_encode, frame_in, writer, test_decode,
+                 &dcodec, &frame_out, &mismatch_seen, &ext_ref);
     frame_in++;
     if (mismatch_seen) break;
   }
@@ -345,6 +374,8 @@
     if (aom_codec_destroy(&dcodec))
       die_codec(&dcodec, "Failed to destroy decoder");
 
+  if (allocated_raw_shift) aom_img_free(&raw_shift);
+  aom_img_free(&ext_ref);
   aom_img_free(&raw);
   if (aom_codec_destroy(&ecodec))
     die_codec(&ecodec, "Failed to destroy encoder.");