Refactor palette functions to write a general color map
This will allow the mrc-tx experiment to pass in its own color map
and set of probabilities to the tokenize and bit packing functions.
The corresponding change on the decoder side will come in a followup
shortly.
No change in performance.
Change-Id: I1872b7f1b1cf5b102e6289388ce1ae449e67cb4d
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index 1582250..0c211b1 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -624,7 +624,7 @@
++p;
--num;
for (int i = 0; i < num; ++i) {
- aom_write_symbol(w, p->token, p->palette_cdf, n);
+ aom_write_symbol(w, p->token, p->color_map_cdf, n);
++p;
}
*tp = p;
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index a84322a..e4df8ab 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -5976,9 +5976,9 @@
for (plane = 0; plane <= 1; ++plane) {
if (mbmi->palette_mode_info.palette_size[plane] > 0) {
if (!dry_run)
- av1_tokenize_palette_sb(x, plane, t, bsize);
+ av1_tokenize_color_map(x, plane, t, bsize, PALETTE_MAP);
else if (dry_run == DRY_RUN_COSTCOEFFS)
- rate += av1_cost_palette_sb(x, plane, bsize);
+ rate += av1_cost_color_map(x, plane, bsize, PALETTE_MAP);
}
}
}
diff --git a/av1/encoder/rdopt.c b/av1/encoder/rdopt.c
index b63c24c..e805137 100644
--- a/av1/encoder/rdopt.c
+++ b/av1/encoder/rdopt.c
@@ -3156,7 +3156,7 @@
color_cache, n_cache,
#endif // CONFIG_PALETTE_DELTA_ENCODING
cpi->common.bit_depth);
- palette_mode_cost += av1_cost_palette_sb(x, 0, bsize);
+ palette_mode_cost += av1_cost_color_map(x, 0, bsize, PALETTE_MAP);
this_model_rd = intra_model_yrd(cpi, x, bsize, palette_mode_cost);
if (*best_model_rd != INT64_MAX &&
this_model_rd > *best_model_rd + (*best_model_rd >> 1))
@@ -5591,7 +5591,7 @@
color_cache, n_cache,
#endif // CONFIG_PALETTE_DELTA_ENCODING
cpi->common.bit_depth);
- this_rate += av1_cost_palette_sb(x, 1, bsize);
+ this_rate += av1_cost_color_map(x, 1, bsize, PALETTE_MAP);
this_rd = RDCOST(x->rdmult, this_rate, tokenonly_rd_stats.dist);
if (this_rd < *best_rd) {
*best_rd = this_rd;
diff --git a/av1/encoder/tokenize.c b/av1/encoder/tokenize.c
index e82263f..ec49099 100644
--- a/av1/encoder/tokenize.c
+++ b/av1/encoder/tokenize.c
@@ -264,6 +264,18 @@
};
#endif
+typedef struct {
+ int rows;
+ int cols;
+ int n_colors;
+ int plane_width;
+ uint8_t *color_map;
+ aom_cdf_prob (
+ *map_cdf)[PALETTE_COLOR_INDEX_CONTEXTS][CDF_SIZE(PALETTE_COLORS)];
+ const int (
+ *color_cost)[PALETTE_SIZES][PALETTE_COLOR_INDEX_CONTEXTS][PALETTE_COLORS];
+} ColorMapParam;
+
#if !CONFIG_PVQ || CONFIG_VAR_TX
static void cost_coeffs_b(int plane, int block, int blk_row, int blk_col,
BLOCK_SIZE plane_bsize, TX_SIZE tx_size, void *arg) {
@@ -329,25 +341,18 @@
}
#endif // !CONFIG_PVQ || CONFIG_VAR_TX
-static int cost_and_tokenize_map_sb(const MACROBLOCK *const x, int plane,
- TOKENEXTRA **t, int calc_rate,
- BLOCK_SIZE bsize) {
- assert(plane == 0 || plane == 1);
- const MACROBLOCKD *const xd = &x->e_mbd;
- const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
- const uint8_t *const color_map = xd->plane[plane].color_index_map;
- const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
+static int cost_and_tokenize_map(ColorMapParam *param, TOKENEXTRA **t,
+ int calc_rate) {
+ const uint8_t *const color_map = param->color_map;
aom_cdf_prob(
- *palette_cdf)[PALETTE_COLOR_INDEX_CONTEXTS][CDF_SIZE(PALETTE_COLORS)] =
- plane ? xd->tile_ctx->palette_uv_color_index_cdf
- : xd->tile_ctx->palette_y_color_index_cdf;
+ *map_cdf)[PALETTE_COLOR_INDEX_CONTEXTS][CDF_SIZE(PALETTE_COLORS)] =
+ param->map_cdf;
const int(*color_cost)[PALETTE_SIZES][PALETTE_COLOR_INDEX_CONTEXTS]
- [PALETTE_COLORS] = plane ? &x->palette_uv_color_cost
- : &x->palette_y_color_cost;
- int plane_block_width, rows, cols;
- const int n = pmi->palette_size[plane];
- av1_get_block_dimensions(bsize, plane, xd, &plane_block_width, NULL, &rows,
- &cols);
+ [PALETTE_COLORS] = param->color_cost;
+ const int plane_block_width = param->plane_width;
+ const int rows = param->rows;
+ const int cols = param->cols;
+ const int n = param->n_colors;
int this_rate = 0;
uint8_t color_order[PALETTE_MAX_SIZE];
@@ -368,7 +373,7 @@
(*color_cost)[n - PALETTE_MIN_SIZE][color_ctx][color_new_idx];
} else {
(*t)->token = color_new_idx;
- (*t)->palette_cdf = palette_cdf[n - PALETTE_MIN_SIZE][color_ctx];
+ (*t)->color_map_cdf = map_cdf[n - PALETTE_MIN_SIZE][color_ctx];
++(*t);
}
}
@@ -377,21 +382,64 @@
return 0;
}
-int av1_cost_palette_sb(const MACROBLOCK *const x, int plane,
- BLOCK_SIZE bsize) {
- return cost_and_tokenize_map_sb(x, plane, NULL, 1, bsize);
+static void get_palette_params(const MACROBLOCK *const x, int plane,
+ BLOCK_SIZE bsize, ColorMapParam *params) {
+ const MACROBLOCKD *const xd = &x->e_mbd;
+ const MB_MODE_INFO *const mbmi = &xd->mi[0]->mbmi;
+ const PALETTE_MODE_INFO *const pmi = &mbmi->palette_mode_info;
+ params->color_map = xd->plane[plane].color_index_map;
+ params->map_cdf = plane ? xd->tile_ctx->palette_uv_color_index_cdf
+ : xd->tile_ctx->palette_y_color_index_cdf;
+ params->color_cost =
+ plane ? &x->palette_uv_color_cost : &x->palette_y_color_cost;
+ params->n_colors = pmi->palette_size[plane];
+ av1_get_block_dimensions(bsize, plane, xd, ¶ms->plane_width, NULL,
+ ¶ms->rows, ¶ms->cols);
}
-void av1_tokenize_palette_sb(const MACROBLOCK *const x, int plane,
- TOKENEXTRA **t, BLOCK_SIZE bsize) {
+#if CONFIG_MRC_TX
+static void get_mrc_params(const MACROBLOCK *const x, int plane,
+ BLOCK_SIZE bsize, ColorMapParam *params) {
+ // TODO(sarahparker)
+ (void)x;
+ (void)plane;
+ (void)bsize;
+ memset(params, 0, sizeof(*params));
+}
+#endif // CONFIG_MRC_TX
+
+static void get_color_map_params(const MACROBLOCK *const x, int plane,
+ BLOCK_SIZE bsize, COLOR_MAP_TYPE type,
+ ColorMapParam *params) {
+ memset(params, 0, sizeof(*params));
+ switch (type) {
+ case PALETTE_MAP: get_palette_params(x, plane, bsize, params); break;
+#if CONFIG_MRC_TX
+ case MRC_MAP: get_mrc_params(x, plane, bsize, params); break;
+#endif // CONFIG_MRC_TX
+ default: assert(0 && "Invalid color map type"); return;
+ }
+}
+
+int av1_cost_color_map(const MACROBLOCK *const x, int plane, BLOCK_SIZE bsize,
+ COLOR_MAP_TYPE type) {
assert(plane == 0 || plane == 1);
- const MACROBLOCKD *const xd = &x->e_mbd;
- const uint8_t *const color_map = xd->plane[plane].color_index_map;
+ ColorMapParam color_map_params;
+ get_color_map_params(x, plane, bsize, type, &color_map_params);
+ return cost_and_tokenize_map(&color_map_params, NULL, 1);
+}
+
+void av1_tokenize_color_map(const MACROBLOCK *const x, int plane,
+ TOKENEXTRA **t, BLOCK_SIZE bsize,
+ COLOR_MAP_TYPE type) {
+ assert(plane == 0 || plane == 1);
+ ColorMapParam color_map_params;
+ get_color_map_params(x, plane, bsize, type, &color_map_params);
// The first color index does not use context or entropy.
- (*t)->token = color_map[0];
- (*t)->palette_cdf = NULL;
+ (*t)->token = color_map_params.color_map[0];
+ (*t)->color_map_cdf = NULL;
++(*t);
- cost_and_tokenize_map_sb(x, plane, t, 0, bsize);
+ cost_and_tokenize_map(&color_map_params, t, 0);
}
#if CONFIG_PVQ
diff --git a/av1/encoder/tokenize.h b/av1/encoder/tokenize.h
index 08b68f2..c8a5948 100644
--- a/av1/encoder/tokenize.h
+++ b/av1/encoder/tokenize.h
@@ -37,7 +37,7 @@
typedef struct {
aom_cdf_prob (*tail_cdf)[CDF_SIZE(ENTROPY_TOKENS)];
aom_cdf_prob (*head_cdf)[CDF_SIZE(ENTROPY_TOKENS)];
- aom_cdf_prob *palette_cdf;
+ aom_cdf_prob *color_map_cdf;
int eob_val;
int first_val;
const aom_prob *context_tree;
@@ -75,10 +75,12 @@
int mi_col, BLOCK_SIZE bsize, int *rate);
#endif
-int av1_cost_palette_sb(const MACROBLOCK *const x, int plane, BLOCK_SIZE bsize);
+int av1_cost_color_map(const MACROBLOCK *const x, int plane, BLOCK_SIZE bsize,
+ COLOR_MAP_TYPE type);
-void av1_tokenize_palette_sb(const MACROBLOCK *const x, int plane,
- TOKENEXTRA **t, BLOCK_SIZE bsize);
+void av1_tokenize_color_map(const MACROBLOCK *const x, int plane,
+ TOKENEXTRA **t, BLOCK_SIZE bsize,
+ COLOR_MAP_TYPE type);
void av1_tokenize_sb(const struct AV1_COMP *cpi, struct ThreadData *td,
TOKENEXTRA **t, RUN_TYPE dry_run, BLOCK_SIZE bsize,
